Unregistering a node should also remove it from tokenMap if it is there and recalculate the placements

Patch by marcuse; reviewed by Sam Tunnicliffe for CASSANDRA-20346

CHANGES.txt

@@ -1,4 +1,5 @@
 5.1
+ * Unregistering a node should also remove it from tokenMap if it is there and recalculate the placements (CASSANDRA-20346)
  * Fix PartitionUpdate.isEmpty deserialization issue to avoid potential EOFException (CASSANDRA-20345)
  * Avoid adding LEFT nodes to tokenMap on upgrade from gossip (CASSANDRA-20344)
  * Allow empty placements when deserializing cluster metadata (CASSANDRA-20343)

src/java/org/apache/cassandra/service/StorageService.java

@@ -1630,7 +1630,7 @@ public class StorageService extends NotificationBroadcasterSupport implements IE
                         throw new RuntimeException("Can't abort bootstrap for " + nodeId + " since it is not bootstrapping");
                     ClusterMetadataService.instance().commit(new CancelInProgressSequence(nodeId));
                 }
-                ClusterMetadataService.instance().commit(new Unregister(nodeId, EnumSet.of(REGISTERED, BOOTSTRAPPING, BOOT_REPLACING)));
+                ClusterMetadataService.instance().commit(new Unregister(nodeId, EnumSet.of(REGISTERED, BOOTSTRAPPING, BOOT_REPLACING), ClusterMetadataService.instance().placementProvider()));
                 break;
             default:
                 throw new RuntimeException("Can't abort bootstrap for node " + nodeId + " since the state is " + nodeState);

src/java/org/apache/cassandra/tcm/CMSOperations.java

@@ -252,7 +252,7 @@ public class CMSOperations implements CMSOperationsMBean
         for (NodeId nodeId : nodeIds)
         {
             logger.info("Unregistering " + nodeId);
-            cms.commit(new Unregister(nodeId, EnumSet.of(NodeState.LEFT)));
+            cms.commit(new Unregister(nodeId, EnumSet.of(NodeState.LEFT), ClusterMetadataService.instance().placementProvider()));
         }
     }
 }

src/java/org/apache/cassandra/tcm/ClusterMetadata.java

@@ -411,7 +411,9 @@ public class ClusterMetadata
 
         public Transformer unregister(NodeId nodeId)
         {
-            directory = directory.without(nodeId);
+            directory = directory.withoutRackAndDC(nodeId).without(nodeId);
+            if (!tokenMap.tokens(nodeId).isEmpty())
+                tokenMap = tokenMap.unassignTokens(nodeId);
             return this;
         }
 

src/java/org/apache/cassandra/tcm/membership/Directory.java

@@ -252,7 +252,11 @@ public class Directory implements MetadataValue<Directory>
     {
         InetAddressAndPort endpoint = peers.get(id);
         Location location = locations.get(id);
+        if (location == null)
+            return this;
         BTreeMultimap<String, InetAddressAndPort> rackEP = (BTreeMultimap<String, InetAddressAndPort>) racksByDC.get(location.datacenter);
+        if (rackEP == null)
+            return this;
         rackEP = rackEP.without(location.rack, endpoint);
         BTreeMap<String, Multimap<String, InetAddressAndPort>> newRacksByDC;
         if (rackEP.isEmpty())

src/java/org/apache/cassandra/tcm/transformations/Register.java

@@ -115,8 +115,7 @@ public class Register implements Transformation
         if (nodeId == null || metadata.directory.peerState(nodeId) == NodeState.LEFT)
         {
             if (nodeId != null)
-                ClusterMetadataService.instance().commit(new Unregister(nodeId, EnumSet.of(NodeState.LEFT)));
+                ClusterMetadataService.instance().commit(new Unregister(nodeId, EnumSet.of(NodeState.LEFT), ClusterMetadataService.instance().placementProvider()));
-
             Register registration = new Register(nodeAddresses, location, nodeVersion);
             nodeId = ClusterMetadataService.instance().commit(registration).directory.peerId(nodeAddresses.broadcastAddress);
         }

src/java/org/apache/cassandra/tcm/transformations/Unregister.java

@@ -31,6 +31,8 @@ import org.apache.cassandra.tcm.ClusterMetadataService;
 import org.apache.cassandra.tcm.Transformation;
 import org.apache.cassandra.tcm.membership.NodeId;
 import org.apache.cassandra.tcm.membership.NodeState;
+import org.apache.cassandra.tcm.ownership.DataPlacements;
+import org.apache.cassandra.tcm.ownership.PlacementProvider;
 import org.apache.cassandra.tcm.sequences.LockedRanges;
 import org.apache.cassandra.tcm.serialization.AsymmetricMetadataSerializer;
 import org.apache.cassandra.tcm.serialization.Version;
@@ -43,11 +45,13 @@ public class Unregister implements Transformation
 
     private final NodeId nodeId;
     private final EnumSet<NodeState> allowedNodeStartStates;
+    private final PlacementProvider placementProvider;
 
-    public Unregister(NodeId nodeId, EnumSet<NodeState> allowedNodeStartStates)
+    public Unregister(NodeId nodeId, EnumSet<NodeState> allowedNodeStartStates, PlacementProvider placementProvider)
     {
         this.nodeId = nodeId;
         this.allowedNodeStartStates = allowedNodeStartStates;
+        this.placementProvider = placementProvider;
     }
 
     @Override
@@ -67,6 +71,15 @@ public class Unregister implements Transformation
             return new Transformation.Rejected(INVALID, "Can't unregister " + nodeId + " - node state is " + startState + " not " + allowedNodeStartStates);
 
         ClusterMetadata.Transformer next = prev.transformer().unregister(nodeId);
+        if (!prev.tokenMap.tokens(nodeId).isEmpty())
+        {
+            ClusterMetadata nextMetadata = next.build().metadata;
+            DataPlacements placements = placementProvider.calculatePlacements(nextMetadata.epoch,
+                                                                              nextMetadata.tokenMap.toRanges(),
+                                                                              nextMetadata,
+                                                                              nextMetadata.schema.getKeyspaces());
+            next = next.with(placements);
+        }
 
         return Transformation.success(next, LockedRanges.AffectedRanges.EMPTY);
     }
@@ -78,7 +91,7 @@ public class Unregister implements Transformation
     public static void unregister(NodeId nodeId)
     {
         ClusterMetadataService.instance()
-                              .commit(new Unregister(nodeId, EnumSet.allOf(NodeState.class)));
+                              .commit(new Unregister(nodeId, EnumSet.allOf(NodeState.class), ClusterMetadataService.instance().placementProvider()));
     }
 
     public String toString()
@@ -113,7 +126,7 @@ public class Unregister implements Transformation
                     states.add(NodeState.valueOf(in.readUTF()));
             }
             NodeId nodeId = NodeId.serializer.deserialize(in, version);
-            return new Unregister(nodeId, version.isAtLeast(Version.V2) ? states : EnumSet.allOf(NodeState.class));
+            return new Unregister(nodeId, version.isAtLeast(Version.V2) ? states : EnumSet.allOf(NodeState.class), ClusterMetadataService.instance().placementProvider());
         }
 
         public long serializedSize(Transformation t, Version version)

test/distributed/org/apache/cassandra/distributed/test/log/RegisterTest.java

@@ -44,6 +44,7 @@ import org.apache.cassandra.tcm.membership.NodeAddresses;
 import org.apache.cassandra.tcm.membership.NodeId;
 import org.apache.cassandra.tcm.membership.NodeState;
 import org.apache.cassandra.tcm.membership.NodeVersion;
+import org.apache.cassandra.tcm.ownership.PlacementProvider;
 import org.apache.cassandra.tcm.sequences.LeaveStreams;
 import org.apache.cassandra.tcm.sequences.UnbootstrapAndLeave;
 import org.apache.cassandra.tcm.serialization.Version;
@@ -75,15 +76,16 @@ public class RegisterTest extends TestBaseImpl
             for (int i : new int[]{ 3, 2 })
             {
                 cluster.get(i).runOnInstance(() -> {
+                    PlacementProvider pp = ClusterMetadataService.instance().placementProvider();
                     ClusterMetadataService.instance().commit(new PrepareLeave(ClusterMetadata.current().myNodeId(),
                                                                               true,
-                                                                              ClusterMetadataService.instance().placementProvider(),
+                                                                              pp,
                                                                               LeaveStreams.Kind.UNBOOTSTRAP));
                     UnbootstrapAndLeave unbootstrapAndLeave = (UnbootstrapAndLeave) ClusterMetadata.current().inProgressSequences.get(ClusterMetadata.current().myNodeId());
                     ClusterMetadataService.instance().commit(unbootstrapAndLeave.startLeave);
                     ClusterMetadataService.instance().commit(unbootstrapAndLeave.midLeave);
                     ClusterMetadataService.instance().commit(unbootstrapAndLeave.finishLeave);
-                    ClusterMetadataService.instance().commit(new Unregister(ClusterMetadata.current().myNodeId(), EnumSet.of(NodeState.LEFT)));
+                    ClusterMetadataService.instance().commit(new Unregister(ClusterMetadata.current().myNodeId(), EnumSet.of(NodeState.LEFT), pp));
                 });
 
                 cluster.get(1).runOnInstance(() -> {
@@ -142,7 +144,7 @@ public class RegisterTest extends TestBaseImpl
                         }
 
                         // If we unregister oldNode, then the ceiling for serialization version will rise
-                        ClusterMetadataService.instance().commit(new Unregister(oldNode, EnumSet.allOf(NodeState.class)));
+                        ClusterMetadataService.instance().commit(new Unregister(oldNode, EnumSet.allOf(NodeState.class), ClusterMetadataService.instance().placementProvider()));
                         assertEquals(ClusterMetadata.current().directory.clusterMinVersion.serializationVersion,
                                      NodeVersion.CURRENT_METADATA_VERSION.asInt());
                         bytes = t.kind().toVersionedBytes(t);

test/unit/org/apache/cassandra/tcm/UnregisterTest.java

@@ -0,0 +1,148 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.cassandra.tcm;
+
+import java.util.EnumSet;
+import java.util.function.BiFunction;
+
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Sets;
+import com.google.common.collect.Streams;
+import org.junit.Test;
+
+import org.apache.cassandra.dht.Murmur3Partitioner;
+import org.apache.cassandra.dht.Token;
+import org.apache.cassandra.locator.InetAddressAndPort;
+import org.apache.cassandra.schema.DistributedMetadataLogKeyspace;
+import org.apache.cassandra.schema.DistributedSchema;
+import org.apache.cassandra.schema.KeyspaceMetadata;
+import org.apache.cassandra.schema.KeyspaceParams;
+import org.apache.cassandra.schema.Keyspaces;
+import org.apache.cassandra.tcm.membership.Directory;
+import org.apache.cassandra.tcm.membership.Location;
+import org.apache.cassandra.tcm.membership.NodeAddresses;
+import org.apache.cassandra.tcm.membership.NodeId;
+import org.apache.cassandra.tcm.membership.NodeState;
+import org.apache.cassandra.tcm.membership.NodeVersion;
+import org.apache.cassandra.tcm.ownership.UniformRangePlacement;
+import org.apache.cassandra.tcm.transformations.Assassinate;
+import org.apache.cassandra.tcm.transformations.Register;
+import org.apache.cassandra.tcm.transformations.Unregister;
+import org.apache.cassandra.tcm.transformations.UnsafeJoin;
+
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNull;
+
+public class UnregisterTest
+{
+    private final KeyspaceMetadata ksm = KeyspaceMetadata.create("ks_nts", KeyspaceParams.nts("dc1", 3, "dc2", 3));
+    private final Keyspaces kss = Keyspaces.of(DistributedMetadataLogKeyspace.initialMetadata(Sets.newHashSet("dc1", "dc2")), ksm);
+    private final DistributedSchema initialSchema = new DistributedSchema(kss);
+
+    @Test
+    public void testBasicUnregister()
+    {
+        unregisterHelper((toUnregister, metadata) -> {
+            metadata = assassinate(toUnregister, metadata);
+            metadata = unregister(toUnregister, metadata);
+            return metadata;
+        });
+    }
+
+    @Test
+    public void badStateUnregister()
+    {
+        unregisterHelper((toUnregister, metadata) -> {
+            metadata = left(toUnregister, metadata);
+            metadata = unregister(toUnregister, metadata);
+            return metadata;
+        });
+    }
+
+    private void unregisterHelper(BiFunction<Integer, ClusterMetadata, ClusterMetadata> f)
+    {
+        ClusterMetadata metadata = new ClusterMetadata(Murmur3Partitioner.instance, Directory.EMPTY, initialSchema);
+
+        int toUnregister = 5;
+        for (int i = 0; i < 10; i++)
+        {
+            // node we're unregistering is only node in dc2 - make sure the dc is gone after unregistration
+            String dc = "dc"+(i == toUnregister ? "2" :"1");
+            metadata = unsafejoin(i, register(i, dc, metadata));
+        }
+        Token t = token(toUnregister, metadata);
+        InetAddressAndPort ep = ep(toUnregister);
+        NodeId nodeId = metadata.directory.peerId(ep(toUnregister));
+        metadata = f.apply(toUnregister, metadata);
+        assertNoTrace(ep, nodeId, t, metadata);
+    }
+
+    private ClusterMetadata left(int i, ClusterMetadata metadata)
+    {
+        NodeId nodeId = metadata.directory.peerId(ep(i));
+        ClusterMetadata.Transformer t = metadata.transformer().withNodeState(nodeId, NodeState.LEFT);
+        return t.build().metadata;
+    }
+
+    private static void assertNoTrace(InetAddressAndPort ep, NodeId nodeId, Token t, ClusterMetadata metadata)
+    {
+        assertNull(metadata.tokenMap.owner(t));
+        assertFalse(metadata.directory.states.containsKey(nodeId));
+        assertFalse(metadata.directory.peerIds().contains(nodeId));
+        assertFalse(metadata.directory.allAddresses().contains(ep));
+        assertFalse(metadata.directory.allJoinedEndpoints().contains(ep));
+        assertFalse(metadata.directory.allDatacenterRacks().containsKey("dc2"));
+        assertFalse(metadata.directory.knownDatacenters().contains("dc2"));
+        metadata.placements.asMap().forEach((params, placement) -> {
+            assertFalse(Streams.concat(placement.writes.endpoints.stream(), placement.reads.endpoints.stream()).anyMatch((fr) -> fr.endpoints().contains(ep)));
+        });
+    }
+
+    private Token token(int i, ClusterMetadata metadata)
+    {
+        NodeId nodeId = metadata.directory.peerId(ep(i));
+        return metadata.tokenMap.tokens(nodeId).iterator().next();
+    }
+
+    private ClusterMetadata assassinate(int i, ClusterMetadata metadata)
+    {
+        return new Assassinate(metadata.directory.peerId(ep(i)), new UniformRangePlacement()).execute(metadata).success().metadata;
+    }
+
+    private ClusterMetadata unregister(int i, ClusterMetadata metadata )
+    {
+        return new Unregister(metadata.directory.peerId(ep(i)), EnumSet.of(NodeState.LEFT), new UniformRangePlacement()).execute(metadata).success().metadata;
+    }
+
+    private ClusterMetadata register(int i, String dc, ClusterMetadata metadata)
+    {
+        return new Register(new NodeAddresses(ep(i)), new Location(dc, "rack1"), NodeVersion.CURRENT).execute(metadata).success().metadata;
+    }
+
+    private ClusterMetadata unsafejoin(int i, ClusterMetadata metadata)
+    {
+        NodeId nodeId = metadata.directory.peerId(ep(i));
+        return new UnsafeJoin(nodeId, ImmutableSet.of(Murmur3Partitioner.instance.getRandomToken()), new UniformRangePlacement()).execute(metadata).success().metadata;
+    }
+
+    private InetAddressAndPort ep(int i)
+    {
+        return InetAddressAndPort.getByNameUnchecked("127.0.0."+i);
+    }
+}