From 471f782925465ffc81e19fcb750b68c4369c0bc4 Mon Sep 17 00:00:00 2001 From: Yuki Morishita Date: Fri, 31 May 2024 14:56:36 +0900 Subject: [PATCH] Add bootstrap options to control sources of streaming patch by Yuki Morishita; reviewed by xxx for CASSANDRA-19733 --- CHANGES.txt | 1 + .../config/CassandraRelevantProperties.java | 9 + .../apache/cassandra/dht/BootStrapper.java | 3 + .../cassandra/dht/BootstrapSourceFilter.java | 224 +++++++++ .../service/BootstrapOptionsParser.java | 90 ++++ .../dht/BootstrapSourceFilterTest.java | 425 ++++++++++++++++++ 6 files changed, 752 insertions(+) create mode 100644 src/java/org/apache/cassandra/dht/BootstrapSourceFilter.java create mode 100644 src/java/org/apache/cassandra/service/BootstrapOptionsParser.java create mode 100644 test/unit/org/apache/cassandra/dht/BootstrapSourceFilterTest.java diff --git a/CHANGES.txt b/CHANGES.txt index 95ee1c9979..3c181eed94 100644 --- a/CHANGES.txt +++ b/CHANGES.txt @@ -70,6 +70,7 @@ Merged from 4.0: * Revert unnecessary read lock acquisition when reading ring version in TokenMetadata introduced in CASSANDRA-16286 (CASSANDRA-19107) Merged from 3.11: * Move ClientWarn.State#warnings to a thread-safe list (CASSANDRA-19427) + * Add bootstrap options to control sources of streaming (CASSANDRA-19733) Merged from 3.0: * Fix SCM URL link (CASSANDRA-19422) * Backport CASSANDRA-16418 to 3.x (CASSANDRA-18824) diff --git a/src/java/org/apache/cassandra/config/CassandraRelevantProperties.java b/src/java/org/apache/cassandra/config/CassandraRelevantProperties.java index 968feae396..6c3616f01a 100644 --- a/src/java/org/apache/cassandra/config/CassandraRelevantProperties.java +++ b/src/java/org/apache/cassandra/config/CassandraRelevantProperties.java @@ -61,6 +61,15 @@ public enum CassandraRelevantProperties AUTO_REPAIR_FREQUENCY_SECONDS("cassandra.auto_repair_frequency_seconds", convertToString(TimeUnit.MINUTES.toSeconds(5))), BATCHLOG_REPLAY_TIMEOUT_IN_MS("cassandra.batchlog.replay_timeout_in_ms"), BATCH_COMMIT_LOG_SYNC_INTERVAL("cassandra.batch_commitlog_sync_interval_millis", "1000"), + /** The system property to specify the sources to exclude in specific dc and rack. */ + BOOTSTRAP_EXCLUDE_DCS("cassandra.bootstrap.exclude_dcs"), + /** The system property to specify the sources in specific dc and rack. */ + BOOTSTRAP_INCLUDE_DCS("cassandra.bootstrap.include_dcs"), + /** + * The system property to specify the sources using their IP addresses. + * Can be used with {@code cassandra.bootstrap.include_dcs} to further restrict the sources. + */ + BOOTSTRAP_INCLUDE_SOURCES( "cassandra.bootstrap.include_sources"), /** * When bootstraping how long to wait for schema versions to be seen. */ diff --git a/src/java/org/apache/cassandra/dht/BootStrapper.java b/src/java/org/apache/cassandra/dht/BootStrapper.java index 5d5529e15f..d782945559 100644 --- a/src/java/org/apache/cassandra/dht/BootStrapper.java +++ b/src/java/org/apache/cassandra/dht/BootStrapper.java @@ -33,6 +33,7 @@ import org.apache.cassandra.gms.Gossiper; import org.apache.cassandra.locator.AbstractReplicationStrategy; import org.apache.cassandra.locator.InetAddressAndPort; import org.apache.cassandra.locator.TokenMetadata; +import org.apache.cassandra.service.BootstrapOptionsParser; import org.apache.cassandra.service.StorageService; import org.apache.cassandra.streaming.*; import org.apache.cassandra.utils.FBUtilities; @@ -73,6 +74,8 @@ public class BootStrapper extends ProgressEventNotifierSupport stateStore, true, DatabaseDescriptor.getStreamingConnectionsPerHost()); + BootstrapSourceFilter filter = BootstrapOptionsParser.parse(tokenMetadata, DatabaseDescriptor.getEndpointSnitch()); + streamer.addSourceFilter(filter); final Collection nonLocalStrategyKeyspaces = Schema.instance.distributedKeyspaces().names(); if (nonLocalStrategyKeyspaces.isEmpty()) logger.debug("Schema does not contain any non-local keyspaces to stream on bootstrap"); diff --git a/src/java/org/apache/cassandra/dht/BootstrapSourceFilter.java b/src/java/org/apache/cassandra/dht/BootstrapSourceFilter.java new file mode 100644 index 0000000000..239ce1a860 --- /dev/null +++ b/src/java/org/apache/cassandra/dht/BootstrapSourceFilter.java @@ -0,0 +1,224 @@ +/* + * 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.dht; + +import java.util.Collection; +import java.util.HashSet; +import java.util.Map; +import java.util.Set; + +import com.google.common.collect.HashMultimap; +import com.google.common.collect.ImmutableMultimap; +import com.google.common.collect.SetMultimap; + +import org.apache.cassandra.locator.IEndpointSnitch; +import org.apache.cassandra.locator.InetAddressAndPort; +import org.apache.cassandra.locator.Replica; +import org.apache.cassandra.locator.TokenMetadata; + +/** + * Defines options for the bootstrap process to restrict the sources from which to stream data. + */ +public class BootstrapSourceFilter implements RangeStreamer.SourceFilter +{ + /* + * Special value to include all racks in a datacenter. + * A whitespace is used to avoid conflict with the actual rack name. + * It is assumed to be safe because the rack name cannot contain just a whitespace, + * since the implementations of IEndpointSnitch are expected to trim the rack name. + */ + private static final String ALL_RACKS = " "; + + private final Set sources; + private final ImmutableMultimap includedDcRacks; + private final ImmutableMultimap excludedDcRacks; + private final IEndpointSnitch snitch; + + private BootstrapSourceFilter(Set sources, + ImmutableMultimap includedDcRacks, + ImmutableMultimap excludedDcRacks, + IEndpointSnitch snitch) + { + this.sources = sources; + this.includedDcRacks = includedDcRacks; + this.excludedDcRacks = excludedDcRacks; + this.snitch = snitch; + } + + @Override + public boolean apply(Replica replica) + { + String dc = snitch.getDatacenter(replica); + String rack = snitch.getRack(replica.endpoint()); + boolean include = sources.isEmpty() || sources.contains(replica.endpoint()); + include &= includedDcRacks.isEmpty() || includedDcRacks.containsEntry(dc, ALL_RACKS) + || includedDcRacks.containsEntry(dc, rack); + include &= excludedDcRacks.isEmpty() || !excludedDcRacks.containsEntry(dc, ALL_RACKS) + && !excludedDcRacks.containsEntry(dc, rack); + return include; + } + + @Override + public String message(Replica replica) + { + return ""; + } + + /** + * @return the builder for {@link BootstrapSourceFilter}. + */ + public static Builder builder(TokenMetadata tmd, IEndpointSnitch snitch) + { + return new Builder(tmd.cloneOnlyTokenMap().getAllMembers(), snitch); + } + + /** + * Builder class for {@link BootstrapSourceFilter}. + */ + public static class Builder + { + private final Set sources = new HashSet<>(); + private final SetMultimap includedDcRacks = HashMultimap.create(); + private final SetMultimap excludedDcRacks = HashMultimap.create(); + private final Collection allMembers; + private final IEndpointSnitch snitch; + + private Builder(Collection allMembers, IEndpointSnitch snitch) + { + this.allMembers = allMembers; + this.snitch = snitch; + } + + /** + * Add source to include in the bootstrap process. + * + * @param source inet address of the source to include + * @return this builder + * @throws IllegalArgumentException if the source is not part of the token metadata + */ + public Builder include(InetAddressAndPort source) + { + if (!allMembers.contains(source)) + throw new IllegalArgumentException(String.format("The source %s is not part of the cluster", source)); + this.sources.add(source); + return this; + } + + /** + * Add srouces from specified datacenters to include in the bootstrap process. + * + * @param datacenter the name of the datacenter to include sources from + * @return this builder + * @throws IllegalArgumentException if the dc is not part of the token metadata + */ + public Builder includeDc(String datacenter) + { + return includeDcRack(datacenter, null); + } + + /** + * Add srouces from specified datacenters and racks to include in the bootstrap process. + * If the same datacenter name or rack name is specified multiple times using this method, the last one is used. + * + * @param datacenter the name of the datacenter + * @param rack the name of the rack in datacenter to include sources from. If null, include all racks in the datacenter. + * @return this builder + * @throws IllegalArgumentException if the dc is not part of the token metadata + */ + public Builder includeDcRack(String datacenter, String rack) + { + validateDcRack(datacenter, rack); + if (rack == null) + includedDcRacks.put(datacenter, ALL_RACKS); + else + includedDcRacks.put(datacenter, rack); + return this; + } + + public Builder excludeDc(String datacenter) + { + return excludeDcRack(datacenter, null); + } + + /** + * Exclude srouces from specified datacenters and racks to include in the bootstrap process. + * + * @param datacenter the name of the datacenter + * @param rack the name of the rack in datacenter to include sources from. If null, include all racks in the datacenter. + * @return this builder + * @throws IllegalArgumentException if the dc is not part of the token metadata + */ + public Builder excludeDcRack(String datacenter, String rack) + { + validateDcRack(datacenter, rack); + if (rack == null) + excludedDcRacks.put(datacenter, ALL_RACKS); + else + excludedDcRacks.put(datacenter, rack); + return this; + } + + public BootstrapSourceFilter build() + { + // Validate conflicting case + // - the same datacenter/rack pair is both included and excluded + for (Map.Entry dcRack : includedDcRacks.entries()) + { + if (excludedDcRacks.containsEntry(dcRack.getKey(), dcRack.getValue())) + throw new IllegalArgumentException(String.format("%s%s is included and excluded", + dcRack.getKey(), + ALL_RACKS.equals(dcRack.getValue()) ? "" : ':' + dcRack.getValue())); + } + + return new BootstrapSourceFilter(sources, + ImmutableMultimap.copyOf(includedDcRacks), + ImmutableMultimap.copyOf(excludedDcRacks), + snitch); + } + + /** + * Validate that given datacenter and rack pair exists in the cluster. + * + * @param datacenter name of the datacenter + * @param rack name of the rack + */ + private void validateDcRack(String datacenter, String rack) + { + if (datacenter == null) + throw new IllegalArgumentException("Datacenter name cannot be null"); + + boolean found = false; + for (InetAddressAndPort node : allMembers) { + if (datacenter.equals(snitch.getDatacenter(node))) { + if (rack == null || rack.equals(snitch.getRack(node))) { + found = true; + break; + } + } + } + + if (!found) { + String message = rack == null + ? String.format("%s is not part of the cluster", datacenter) + : String.format("%s:%s is not part of the cluster", datacenter, rack); + throw new IllegalArgumentException(message); + } + } + } +} \ No newline at end of file diff --git a/src/java/org/apache/cassandra/service/BootstrapOptionsParser.java b/src/java/org/apache/cassandra/service/BootstrapOptionsParser.java new file mode 100644 index 0000000000..834a4fc7d7 --- /dev/null +++ b/src/java/org/apache/cassandra/service/BootstrapOptionsParser.java @@ -0,0 +1,90 @@ +/* + * 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.service; + + +import org.apache.cassandra.dht.BootstrapSourceFilter; +import org.apache.cassandra.locator.IEndpointSnitch; +import org.apache.cassandra.locator.InetAddressAndPort; +import org.apache.cassandra.locator.TokenMetadata; + +import static org.apache.cassandra.config.CassandraRelevantProperties.*; + +/** + * Constructs {@link BootstrapSourceFilter} from system properties and others. + */ +public class BootstrapOptionsParser +{ + + public static BootstrapSourceFilter parse(TokenMetadata tmd, IEndpointSnitch snitch) + { + BootstrapSourceFilter.Builder builder = BootstrapSourceFilter.builder(tmd, snitch); + + String includeDcRackOption = BOOTSTRAP_INCLUDE_DCS.getString(); + if (includeDcRackOption != null) + { + String[] dcRacks = includeDcRackOption.split(","); + for (String dcRack : dcRacks) + { + String[] parts = dcRack.split(":"); + if (parts.length == 1) + builder.includeDc(parts[0].trim()); + else if (parts.length == 2) + builder.includeDcRack(parts[0].trim(), parts[1].trim()); + else + throw new IllegalArgumentException("Invalid dc:rack option: " + includeDcRackOption); + } + } + String excludeDcRackOption = BOOTSTRAP_EXCLUDE_DCS.getString(); + if (excludeDcRackOption != null) + { + String[] dcRacks = excludeDcRackOption.split(","); + for (String dcRack : dcRacks) + { + String[] parts = dcRack.split(":"); + if (parts.length == 1) + builder.excludeDc(parts[0].trim()); + else if (parts.length == 2) + builder.excludeDcRack(parts[0].trim(), parts[1].trim()); + else + throw new IllegalArgumentException("Invalid dc:rack option: " + excludeDcRackOption); + } + } + + String includeSourcesOption = BOOTSTRAP_INCLUDE_SOURCES.getString(); + if (includeSourcesOption != null) + { + String[] sources = includeSourcesOption.split(","); + for (String source : sources) + { + try + { + InetAddressAndPort sourceAddress = InetAddressAndPort.getByName(source.trim()); + builder.include(sourceAddress); + } + catch (Exception e) + { + throw new IllegalArgumentException("Invalid source address in option: " + includeSourcesOption, e); + } + } + } + + return builder.build(); + } +} diff --git a/test/unit/org/apache/cassandra/dht/BootstrapSourceFilterTest.java b/test/unit/org/apache/cassandra/dht/BootstrapSourceFilterTest.java new file mode 100644 index 0000000000..07a9101cee --- /dev/null +++ b/test/unit/org/apache/cassandra/dht/BootstrapSourceFilterTest.java @@ -0,0 +1,425 @@ +/* + * 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.dht; + +import java.net.UnknownHostException; +import java.util.Collections; +import java.util.Set; +import java.util.UUID; +import java.util.stream.Collectors; +import java.util.stream.StreamSupport; + +import com.google.common.collect.ImmutableMap; +import org.junit.BeforeClass; +import org.junit.Test; + +import org.apache.cassandra.config.DatabaseDescriptor; +import org.apache.cassandra.locator.AbstractNetworkTopologySnitch; +import org.apache.cassandra.locator.AbstractReplicationStrategy; +import org.apache.cassandra.locator.IEndpointSnitch; +import org.apache.cassandra.locator.InetAddressAndPort; +import org.apache.cassandra.locator.Replica; +import org.apache.cassandra.locator.SimpleStrategy; +import org.apache.cassandra.locator.TokenMetadata; +import org.apache.cassandra.service.BootstrapOptionsParser; + +import static org.apache.cassandra.config.CassandraRelevantProperties.*; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +public class BootstrapSourceFilterTest +{ + private static final IEndpointSnitch snitch = new AbstractNetworkTopologySnitch() + { + @Override + public String getRack(InetAddressAndPort endpoint) + { + return "Rack-" + endpoint.addressBytes[2]; + } + + @Override + public String getDatacenter(InetAddressAndPort endpoint) + { + return "DC-" + endpoint.addressBytes[1]; + } + }; + + @BeforeClass + public static void setUpSnitch() + { + DatabaseDescriptor.setEndpointSnitch(snitch); + } + + /** + * Topology that contains 3 DCs, with 2 racks each, and 3 nodes per rack (6 nodes per DC). + */ + private final TokenMetadata tmd; + private final AbstractReplicationStrategy replicationStrategy; + + public BootstrapSourceFilterTest() + { + this.tmd = new TokenMetadata(); + replicationStrategy = new SimpleStrategy("ks", tmd, snitch, ImmutableMap.of(SimpleStrategy.REPLICATION_FACTOR, "1")); + IPartitioner p = new Murmur3Partitioner(); + + // add 2 DCs, with 3 racks each, and 2 nodes per rack (6 nodes per DC) + for (int dc = 1; dc <= 2; dc++) + { + for (int rack = 1; rack <= 3; rack++) + { + for (int node = 1; node <= 2; node++) + { + InetAddressAndPort addr = addr(dc, rack, node); + tmd.updateNormalTokens(Collections.singleton(p.getRandomToken()), addr); + tmd.updateHostId(UUID.randomUUID(), addr); + } + } + } + } + + @Test + public void testEmpty() + { + BootstrapSourceFilter filter = BootstrapSourceFilter.builder(tmd, snitch).build(); + + Set restrictedSources = StreamSupport.stream(replicationStrategy.getRangeAddresses(tmd) + .flattenValues().spliterator(), false) + .filter(filter::apply) + .map(Replica::endpoint) + .collect(Collectors.toSet()); + + // All 12 nodes should be included + assertEquals(12, restrictedSources.size()); + } + + @Test + public void testIncludeSources() + { + BootstrapSourceFilter filter = BootstrapSourceFilter.builder(tmd, snitch) + .include(addr(1, 1, 1)) + .include(addr(1, 1, 2)) + .build(); + + Set restrictedSources = StreamSupport.stream(replicationStrategy.getRangeAddresses(tmd) + .flattenValues().spliterator(), false) + .filter(filter::apply) + .map(Replica::endpoint) + .collect(Collectors.toSet()); + + assertEquals(2, restrictedSources.size()); + assertTrue(restrictedSources.contains(addr(1, 1, 1))); + assertTrue(restrictedSources.contains(addr(1, 1, 2))); + } + + @Test + public void testIncludeDc() + { + BootstrapSourceFilter filter = BootstrapSourceFilter.builder(tmd, snitch) + .includeDc("DC-1") + .build(); + + Set restrictedSources = StreamSupport.stream(replicationStrategy.getRangeAddresses(tmd) + .flattenValues().spliterator(), false) + .filter(filter::apply) + .map(Replica::endpoint) + .collect(Collectors.toSet()); + + // 6 nodes in DC-1 should be included + assertEquals(6, restrictedSources.size()); + assertTrue(restrictedSources.contains(addr(1, 1, 1))); + assertTrue(restrictedSources.contains(addr(1, 1, 2))); + assertTrue(restrictedSources.contains(addr(1, 2, 1))); + assertTrue(restrictedSources.contains(addr(1, 2, 2))); + assertTrue(restrictedSources.contains(addr(1, 3, 1))); + assertTrue(restrictedSources.contains(addr(1, 3, 2))); + } + + @Test + public void testIncludeDcTwice() + { + // When including both the entire DC and the specific rack in DC, the entire DC should be included + BootstrapSourceFilter filter = BootstrapSourceFilter.builder(tmd, snitch) + .includeDcRack("DC-1", "Rack-1") + .includeDc("DC-1") + .build(); + + Set restrictedSources = StreamSupport.stream(replicationStrategy.getRangeAddresses(tmd) + .flattenValues().spliterator(), false) + .filter(filter::apply) + .map(Replica::endpoint) + .collect(Collectors.toSet()); + + // All 6 nodes in DC-1 should be included + assertEquals(6, restrictedSources.size()); + assertTrue(restrictedSources.contains(addr(1, 1, 1))); + assertTrue(restrictedSources.contains(addr(1, 1, 2))); + assertTrue(restrictedSources.contains(addr(1, 2, 1))); + assertTrue(restrictedSources.contains(addr(1, 2, 2))); + assertTrue(restrictedSources.contains(addr(1, 3, 1))); + assertTrue(restrictedSources.contains(addr(1, 3, 2))); + } + + @Test + public void testIncludeDcRack() + { + BootstrapSourceFilter filter = BootstrapSourceFilter.builder(tmd, snitch) + .includeDcRack("DC-1", "Rack-2") + .includeDcRack("DC-2", "Rack-1") + .build(); + + Set restrictedSources = StreamSupport.stream(replicationStrategy.getRangeAddresses(tmd) + .flattenValues().spliterator(), false) + .filter(filter::apply) + .map(Replica::endpoint) + .collect(Collectors.toSet()); + + // 4 nodes from DC-1/Rack-2 and DC-2/Rack-1 should be included + assertEquals(4, restrictedSources.size()); + assertTrue(restrictedSources.contains(addr(1, 2, 1))); + assertTrue(restrictedSources.contains(addr(1, 2, 2))); + assertTrue(restrictedSources.contains(addr(2, 1, 1))); + assertTrue(restrictedSources.contains(addr(2, 1, 2))); + } + + @Test + public void testExcludeDc() + { + BootstrapSourceFilter filter = BootstrapSourceFilter.builder(tmd, snitch) + .excludeDc("DC-1") + .build(); + + Set restrictedSources = StreamSupport.stream(replicationStrategy.getRangeAddresses(tmd) + .flattenValues().spliterator(), false) + .filter(filter::apply) + .map(Replica::endpoint) + .collect(Collectors.toSet()); + + // 6 nodes in DC-1 should be included + assertEquals(6, restrictedSources.size()); + assertTrue(restrictedSources.contains(addr(2, 1, 1))); + assertTrue(restrictedSources.contains(addr(2, 1, 2))); + assertTrue(restrictedSources.contains(addr(2, 2, 1))); + assertTrue(restrictedSources.contains(addr(2, 2, 2))); + assertTrue(restrictedSources.contains(addr(2, 3, 1))); + assertTrue(restrictedSources.contains(addr(2, 3, 2))); + } + + @Test + public void testExcludeDcRack() + { + BootstrapSourceFilter filter = BootstrapSourceFilter.builder(tmd, snitch) + .excludeDcRack("DC-1", "Rack-2") + .excludeDc("DC-2") + .build(); + + Set restrictedSources = StreamSupport.stream(replicationStrategy.getRangeAddresses(tmd) + .flattenValues().spliterator(), false) + .filter(filter::apply) + .map(Replica::endpoint) + .collect(Collectors.toSet()); + + // 2 nodes in DC-1/Rack-1 and 2 nodes in DC-1/Rack-3 should be included + assertEquals(4, restrictedSources.size()); + assertTrue(restrictedSources.contains(addr(1, 1, 1))); + assertTrue(restrictedSources.contains(addr(1, 1, 2))); + assertTrue(restrictedSources.contains(addr(1, 3, 1))); + assertTrue(restrictedSources.contains(addr(1, 3, 2))); + } + + @Test + public void testIncludeAndExcludeDc() + { + BootstrapSourceFilter filter = BootstrapSourceFilter.builder(tmd, snitch) + .includeDc("DC-1") + .excludeDcRack("DC-1", "Rack-2") + .build(); + + Set restrictedSources = StreamSupport.stream(replicationStrategy.getRangeAddresses(tmd) + .flattenValues().spliterator(), false) + .filter(filter::apply) + .map(Replica::endpoint) + .collect(Collectors.toSet()); + + // 2 nodes in DC-1/Rack-1 and 2 nodes in DC-1/Rack-3 should be included + assertEquals(4, restrictedSources.size()); + assertTrue(restrictedSources.contains(addr(1, 1, 1))); + assertTrue(restrictedSources.contains(addr(1, 1, 2))); + assertTrue(restrictedSources.contains(addr(1, 3, 1))); + assertTrue(restrictedSources.contains(addr(1, 3, 2))); + } + + @Test(expected = IllegalArgumentException.class) + public void testIncludeNullSource() + { + BootstrapSourceFilter.builder(tmd, snitch) + .include(null) + .build(); + } + + + @Test(expected = IllegalArgumentException.class) + public void testIncludeNullDc() + { + BootstrapSourceFilter.builder(tmd, snitch) + .includeDc(null) + .build(); + } + + @Test(expected = IllegalArgumentException.class) + public void testIncludeNonexistingDc() + { + BootstrapSourceFilter.builder(tmd, snitch) + .includeDc("Illegal-DC") + .build(); + } + + @Test(expected = IllegalArgumentException.class) + public void testIncludeNonexistingRack() + { + BootstrapSourceFilter.builder(tmd, snitch) + .includeDcRack("DC-1", "Illegal-Rack") + .build(); + } + + @Test(expected = IllegalArgumentException.class) + public void testSameDcIsIncludedAndExcluded() + { + BootstrapSourceFilter.builder(tmd, snitch) + .includeDc("DC-1") + .excludeDc("DC-1") + .build(); + } + + @Test(expected = IllegalArgumentException.class) + public void testSameDcRackIsIncludedAndExcluded() + { + BootstrapSourceFilter.builder(tmd, snitch) + .includeDcRack("DC-1", "Rack-1") + .excludeDcRack("DC-1", "Rack-1") + .build(); + } + + @Test + public void testParseIncludeDcs() + { + BOOTSTRAP_INCLUDE_DCS.setString("DC-1:Rack-1,DC-2"); + try + { + BootstrapSourceFilter filter = BootstrapOptionsParser.parse(tmd, snitch); + + Set restrictedSources = StreamSupport.stream(replicationStrategy.getRangeAddresses(tmd) + .flattenValues().spliterator(), false) + .filter(filter::apply) + .map(Replica::endpoint) + .collect(Collectors.toSet()); + + // 2 nodes in DC-1/Rack-1 and 6 nodes in DC-2 should be included + assertEquals(8, restrictedSources.size()); + assertTrue(restrictedSources.contains(addr(1, 1, 1))); + assertTrue(restrictedSources.contains(addr(1, 1, 2))); + assertTrue(restrictedSources.contains(addr(2, 1, 1))); + assertTrue(restrictedSources.contains(addr(2, 1, 2))); + assertTrue(restrictedSources.contains(addr(2, 2, 1))); + assertTrue(restrictedSources.contains(addr(2, 2, 2))); + assertTrue(restrictedSources.contains(addr(2, 3, 1))); + assertTrue(restrictedSources.contains(addr(2, 3, 2))); + } + finally + { + BOOTSTRAP_INCLUDE_DCS.clearValue(); + } + } + + @Test + public void testParseExcludeDcs() + { + BOOTSTRAP_EXCLUDE_DCS.setString("DC-1:Rack-2,DC-2"); + try + { + BootstrapSourceFilter filter = BootstrapOptionsParser.parse(tmd, snitch); + + Set restrictedSources = StreamSupport.stream(replicationStrategy.getRangeAddresses(tmd) + .flattenValues().spliterator(), false) + .filter(filter::apply) + .map(Replica::endpoint) + .collect(Collectors.toSet()); + + // 2 nodes in DC-1/Rack-1 and 2 nodes in DC-1/Rack-3 should be included + assertEquals(4, restrictedSources.size()); + assertTrue(restrictedSources.contains(addr(1, 1, 1))); + assertTrue(restrictedSources.contains(addr(1, 1, 2))); + assertTrue(restrictedSources.contains(addr(1, 3, 1))); + assertTrue(restrictedSources.contains(addr(1, 3, 2))); + } + finally + { + BOOTSTRAP_EXCLUDE_DCS.clearValue(); + } + } + + @Test + public void testParseIncludeSources() + { + BOOTSTRAP_INCLUDE_SOURCES.setString("1.1.1.1, 1.1.1.2"); + try + { + BootstrapSourceFilter filter = BootstrapOptionsParser.parse(tmd, snitch); + + Set restrictedSources = StreamSupport.stream(replicationStrategy.getRangeAddresses(tmd) + .flattenValues().spliterator(), false) + .filter(filter::apply) + .map(Replica::endpoint) + .collect(Collectors.toSet()); + + assertEquals(2, restrictedSources.size()); + assertTrue(restrictedSources.contains(addr(1, 1, 1))); + assertTrue(restrictedSources.contains(addr(1, 1, 2))); + } + finally + { + BOOTSTRAP_INCLUDE_SOURCES.clearValue(); + } + } + + @Test(expected = IllegalArgumentException.class) + public void testParseIncludeInvalidSources() + { + BOOTSTRAP_INCLUDE_SOURCES.setString("1.1.1.1,invalid1,invalid2"); + try + { + BootstrapOptionsParser.parse(tmd, snitch); + } + finally + { + BOOTSTRAP_INCLUDE_SOURCES.clearValue(); + } + } + + private static InetAddressAndPort addr(int dc, int rack, int node) + { + byte[] addr = new byte[]{ 1, (byte) dc, (byte) rack, (byte) node }; + try + { + return InetAddressAndPort.getByAddress(addr); + } + catch (UnknownHostException e) + { + throw new RuntimeException(e); + } + } +}