Skip to content

Commit 200137b

Browse files
BouncheckBouncheck
authored
Improve DefaultMetadata's TabletMap tablet invalidation logic (#388)
Previously the tablet map would be completely emptied on node removal, necessitating full rebuild of it. With this change the TabletMap will be scanned through on node removals (`RemoveNodeRefresh`) in order to delete only those tablets, that contain removed node as one of the replicas. Additionally we introduce `TabletMapSchemaChangeListener` which will be automatically registered on context initialization as long as schema metadata is enabled. It won't be added if the schema metadata is enabled later at runtime. This listener ensures that relevant tablets will be removed on removals and updates of both keyspaces and tables. The `TabletMapSchemaChangesIT` tests its behaviour. Addresses #377.
1 parent a010b28 commit 200137b

File tree

6 files changed

+273
-3
lines changed

6 files changed

+273
-3
lines changed

‎core/src/main/java/com/datastax/oss/driver/api/core/metadata/TabletMap.java

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,4 +34,25 @@ public interface TabletMap {
3434
* present.
3535
*/
3636
public Tablet getTablet(CqlIdentifier keyspace, CqlIdentifier table, long token);
37+
38+
/**
39+
* Removes all tablets that contain given node in its replica list.
40+
*
41+
* @param node node serving as filter criterion
42+
*/
43+
public void removeByNode(Node node);
44+
45+
/**
46+
* Removes all mappings for a given keyspace.
47+
*
48+
* @param keyspace keyspace to remove
49+
*/
50+
public void removeByKeyspace(CqlIdentifier keyspace);
51+
52+
/**
53+
* Removes all mappings for a given table.
54+
*
55+
* @param table table to remove
56+
*/
57+
public void removeByTable(CqlIdentifier table);
3758
}

‎core/src/main/java/com/datastax/oss/driver/internal/core/context/DefaultDriverContext.java

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,7 @@
6565
import com.datastax.oss.driver.internal.core.metadata.TopologyMonitor;
6666
import com.datastax.oss.driver.internal.core.metadata.schema.MultiplexingSchemaChangeListener;
6767
import com.datastax.oss.driver.internal.core.metadata.schema.NoopSchemaChangeListener;
68+
import com.datastax.oss.driver.internal.core.metadata.schema.TabletMapSchemaChangeListener;
6869
import com.datastax.oss.driver.internal.core.metadata.schema.parsing.DefaultSchemaParserFactory;
6970
import com.datastax.oss.driver.internal.core.metadata.schema.parsing.SchemaParserFactory;
7071
import com.datastax.oss.driver.internal.core.metadata.schema.queries.DefaultSchemaQueriesFactory;
@@ -659,6 +660,10 @@ protected SchemaChangeListener buildSchemaChangeListener(
659660
"com.datastax.oss.driver.internal.core.metadata.schema")
660661
.ifPresent(listeners::add);
661662
}
663+
if (getMetadataManager().isSchemaEnabled()) {
664+
listeners.add(
665+
new TabletMapSchemaChangeListener(getMetadataManager().getMetadata().getTabletMap()));
666+
}
662667
if (listeners.isEmpty()) {
663668
return new NoopSchemaChangeListener(this);
664669
} else if (listeners.size() == 1) {

‎core/src/main/java/com/datastax/oss/driver/internal/core/metadata/DefaultTabletMap.java

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,37 @@ public void addTablet(CqlIdentifier keyspace, CqlIdentifier table, Tablet tablet
122122
existingTablets.add(tablet);
123123
}
124124

125+
/**
126+
* Removes tablets that contain given node (by reference equality) as one of the replicas.
127+
*
128+
* @param node node used as a filter criterion.
129+
*/
130+
@Override
131+
public void removeByNode(Node node) {
132+
for (ConcurrentSkipListSet<Tablet> tabletSet : mapping.values()) {
133+
Iterator<Tablet> it = tabletSet.iterator();
134+
while (it.hasNext()) {
135+
if (it.next().getReplicaNodes().contains(node)) {
136+
it.remove();
137+
}
138+
}
139+
}
140+
}
141+
142+
@Override
143+
public void removeByKeyspace(CqlIdentifier keyspace) {
144+
mapping
145+
.keySet()
146+
.removeIf(keyspaceTableNamePair -> keyspaceTableNamePair.getKeyspace().equals(keyspace));
147+
}
148+
149+
@Override
150+
public void removeByTable(CqlIdentifier table) {
151+
mapping
152+
.keySet()
153+
.removeIf(keyspaceTableNamePair -> keyspaceTableNamePair.getTableName().equals(table));
154+
}
155+
125156
/**
126157
* Represents a single tablet created from tablets-routing-v1 custom payload. Its {@code
127158
* compareTo} implementation intentionally relies solely on {@code lastToken} in order to allow

‎core/src/main/java/com/datastax/oss/driver/internal/core/metadata/RemoveNodeRefresh.java

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -67,10 +67,9 @@ public Result compute(
6767
return new Result(oldMetadata);
6868
} else {
6969
LOG.debug("[{}] Removing node {}", logPrefix, removedNode);
70-
LOG.debug("[{}] Tablet metadata will be wiped and rebuilt due to node removal.", logPrefix);
71-
DefaultMetadata newerMetadata = oldMetadata.withTabletMap(DefaultTabletMap.emptyMap());
70+
oldMetadata.tabletMap.removeByNode(removedNode);
7271
return new Result(
73-
newerMetadata.withNodes(newNodesBuilder.build(), tokenMapEnabled, false, null, context),
72+
oldMetadata.withNodes(newNodesBuilder.build(), tokenMapEnabled, false, null, context),
7473
ImmutableList.of(NodeStateEvent.removed((DefaultNode) removedNode)));
7574
}
7675
}

‎core/src/main/java/com/datastax/oss/driver/internal/core/metadata/schema/TabletMapSchemaChangeListener.java

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
package com.datastax.oss.driver.internal.core.metadata.schema;
2+
3+
import com.datastax.oss.driver.api.core.metadata.TabletMap;
4+
import com.datastax.oss.driver.api.core.metadata.schema.KeyspaceMetadata;
5+
import com.datastax.oss.driver.api.core.metadata.schema.SchemaChangeListenerBase;
6+
import com.datastax.oss.driver.api.core.metadata.schema.TableMetadata;
7+
import edu.umd.cs.findbugs.annotations.NonNull;
8+
9+
public class TabletMapSchemaChangeListener extends SchemaChangeListenerBase {
10+
private final TabletMap tabletMap;
11+
12+
public TabletMapSchemaChangeListener(TabletMap tabletMap) {
13+
this.tabletMap = tabletMap;
14+
}
15+
16+
@Override
17+
public void onKeyspaceDropped(@NonNull KeyspaceMetadata keyspace) {
18+
tabletMap.removeByKeyspace(keyspace.getName());
19+
}
20+
21+
@Override
22+
public void onKeyspaceUpdated(
23+
@NonNull KeyspaceMetadata current, @NonNull KeyspaceMetadata previous) {
24+
tabletMap.removeByKeyspace(previous.getName());
25+
}
26+
27+
@Override
28+
public void onTableDropped(@NonNull TableMetadata table) {
29+
tabletMap.removeByTable(table.getName());
30+
}
31+
32+
@Override
33+
public void onTableUpdated(@NonNull TableMetadata current, @NonNull TableMetadata previous) {
34+
tabletMap.removeByTable(previous.getName());
35+
}
36+
}

‎integration-tests/src/test/java/com/datastax/oss/driver/core/metadata/TabletMapSchemaChangesIT.java

Lines changed: 178 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,178 @@
1+
package com.datastax.oss.driver.core.metadata;
2+
3+
import static org.assertj.core.api.Assertions.assertThat;
4+
import static org.awaitility.Awaitility.await;
5+
6+
import com.datastax.oss.driver.api.core.CqlIdentifier;
7+
import com.datastax.oss.driver.api.core.CqlSession;
8+
import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
9+
import com.datastax.oss.driver.api.core.cql.BoundStatement;
10+
import com.datastax.oss.driver.api.core.cql.PreparedStatement;
11+
import com.datastax.oss.driver.api.core.metadata.KeyspaceTableNamePair;
12+
import com.datastax.oss.driver.api.core.metadata.Node;
13+
import com.datastax.oss.driver.api.core.metadata.schema.KeyspaceMetadata;
14+
import com.datastax.oss.driver.api.core.metadata.schema.TableMetadata;
15+
import com.datastax.oss.driver.api.testinfra.CassandraSkip;
16+
import com.datastax.oss.driver.api.testinfra.ScyllaRequirement;
17+
import com.datastax.oss.driver.api.testinfra.ccm.CustomCcmRule;
18+
import com.datastax.oss.driver.api.testinfra.session.SessionRule;
19+
import com.datastax.oss.driver.api.testinfra.session.SessionUtils;
20+
import com.datastax.oss.driver.internal.core.loadbalancing.BasicLoadBalancingPolicy;
21+
import com.datastax.oss.driver.internal.core.metadata.schema.TabletMapSchemaChangeListener;
22+
import java.time.Duration;
23+
import java.util.concurrent.TimeUnit;
24+
import org.junit.Before;
25+
import org.junit.ClassRule;
26+
import org.junit.Test;
27+
import org.junit.rules.RuleChain;
28+
import org.junit.rules.TestRule;
29+
import org.mockito.ArgumentCaptor;
30+
import org.mockito.Mockito;
31+
32+
@ScyllaRequirement(
33+
minOSS = "6.0.0",
34+
minEnterprise = "2024.2",
35+
description = "Needs to support tablets")
36+
@CassandraSkip(description = "Tablets are ScyllaDB-only extension")
37+
// Ensures that TabletMap used by MetadataManager behaves as desired on certain events
38+
public class TabletMapSchemaChangesIT {
39+
40+
// Same listener as the one registered on initialization by
41+
// DefaultDriverContext#buildSchemaChangeListener
42+
// for TabletMap updates. Note that this mock only verifies that it reacts to ".onXhappening()"
43+
// calls and the
44+
// actual working listener updates the TabletMap.
45+
private static final TabletMapSchemaChangeListener listener =
46+
Mockito.mock(TabletMapSchemaChangeListener.class);
47+
private static final CustomCcmRule CCM_RULE =
48+
CustomCcmRule.builder()
49+
.withNodes(2)
50+
.withCassandraConfiguration(
51+
"experimental_features", "['consistent-topology-changes','tablets']")
52+
.build();
53+
private static final SessionRule<CqlSession> SESSION_RULE =
54+
SessionRule.builder(CCM_RULE)
55+
.withConfigLoader(
56+
SessionUtils.configLoaderBuilder()
57+
.withDuration(DefaultDriverOption.REQUEST_TIMEOUT, Duration.ofSeconds(15))
58+
.withClass(
59+
DefaultDriverOption.LOAD_BALANCING_POLICY_CLASS,
60+
BasicLoadBalancingPolicy.class)
61+
.withBoolean(DefaultDriverOption.METADATA_SCHEMA_ENABLED, true)
62+
.withDuration(DefaultDriverOption.METADATA_SCHEMA_WINDOW, Duration.ofSeconds(0))
63+
.build())
64+
.withSchemaChangeListener(listener)
65+
.build();
66+
67+
@ClassRule
68+
public static final TestRule CHAIN = RuleChain.outerRule(CCM_RULE).around(SESSION_RULE);
69+
70+
private static final int INITIAL_TABLETS = 32;
71+
private static final int REPLICATION_FACTOR = 1;
72+
private static final String KEYSPACE_NAME = "TabletMapSchemaChangesIT";
73+
private static final String TABLE_NAME = "testTable";
74+
private static final KeyspaceTableNamePair TABLET_MAP_KEY =
75+
new KeyspaceTableNamePair(
76+
CqlIdentifier.fromCql(KEYSPACE_NAME), CqlIdentifier.fromCql(TABLE_NAME));
77+
private static final String CREATE_KEYSPACE_QUERY =
78+
"CREATE KEYSPACE IF NOT EXISTS "
79+
+ KEYSPACE_NAME
80+
+ " WITH replication = {'class': "
81+
+ "'NetworkTopologyStrategy', "
82+
+ "'replication_factor': '"
83+
+ REPLICATION_FACTOR
84+
+ "'} AND durable_writes = true AND tablets = "
85+
+ "{'initial': "
86+
+ INITIAL_TABLETS
87+
+ "};";
88+
private static final String CREATE_TABLE_QUERY =
89+
"CREATE TABLE IF NOT EXISTS "
90+
+ KEYSPACE_NAME
91+
+ "."
92+
+ TABLE_NAME
93+
+ " (pk int, ck int, PRIMARY KEY(pk, ck));";
94+
private static final String DROP_KEYSPACE = "DROP KEYSPACE IF EXISTS " + KEYSPACE_NAME;
95+
96+
private static final String INSERT_QUERY_TEMPLATE =
97+
"INSERT INTO " + KEYSPACE_NAME + "." + TABLE_NAME + " (pk, ck) VALUES (%s, %s)";
98+
private static final String SELECT_QUERY_TEMPLATE =
99+
"SELECT pk, ck FROM " + KEYSPACE_NAME + "." + TABLE_NAME + " WHERE pk = ?";
100+
101+
private static final long NOTIF_TIMEOUT_MS = TimeUnit.MINUTES.toMillis(1);
102+
103+
@Before
104+
public void setup() {
105+
SESSION_RULE.session().execute(DROP_KEYSPACE);
106+
SESSION_RULE.session().execute(CREATE_KEYSPACE_QUERY);
107+
SESSION_RULE.session().execute(CREATE_TABLE_QUERY);
108+
SESSION_RULE.session().execute(String.format(INSERT_QUERY_TEMPLATE, "1", "2"));
109+
SESSION_RULE.session().execute(String.format(INSERT_QUERY_TEMPLATE, "3", "4"));
110+
PreparedStatement ps = SESSION_RULE.session().prepare(SELECT_QUERY_TEMPLATE);
111+
BoundStatement bs = ps.bind(1);
112+
// This ensures we hit the node that is not tablet replica
113+
for (Node node : SESSION_RULE.session().getMetadata().getNodes().values()) {
114+
SESSION_RULE.session().execute(bs.setNode(node));
115+
}
116+
// Make sure the tablet information is present
117+
await()
118+
.atMost(30, TimeUnit.SECONDS)
119+
.until(
120+
() ->
121+
SESSION_RULE
122+
.session()
123+
.getMetadata()
124+
.getTabletMap()
125+
.getMapping()
126+
.containsKey(TABLET_MAP_KEY));
127+
// Reset invocations for the next test method
128+
Mockito.clearInvocations(listener);
129+
}
130+
131+
@Test
132+
public void should_remove_tablets_on_keyspace_update() {
133+
SESSION_RULE
134+
.session()
135+
.execute("ALTER KEYSPACE " + KEYSPACE_NAME + " WITH durable_writes = false");
136+
ArgumentCaptor<KeyspaceMetadata> previous = ArgumentCaptor.forClass(KeyspaceMetadata.class);
137+
Mockito.verify(listener, Mockito.timeout(NOTIF_TIMEOUT_MS).times(1))
138+
.onKeyspaceUpdated(Mockito.any(), previous.capture());
139+
assertThat(previous.getValue().getName()).isEqualTo(CqlIdentifier.fromCql(KEYSPACE_NAME));
140+
assertThat(SESSION_RULE.session().getMetadata().getTabletMap().getMapping().keySet())
141+
.doesNotContain(TABLET_MAP_KEY);
142+
}
143+
144+
@Test
145+
public void should_remove_tablets_on_keyspace_drop() {
146+
SESSION_RULE.session().execute(DROP_KEYSPACE);
147+
ArgumentCaptor<KeyspaceMetadata> keyspace = ArgumentCaptor.forClass(KeyspaceMetadata.class);
148+
Mockito.verify(listener, Mockito.timeout(NOTIF_TIMEOUT_MS).times(1))
149+
.onKeyspaceDropped(keyspace.capture());
150+
assertThat(keyspace.getValue().getName()).isEqualTo(CqlIdentifier.fromCql(KEYSPACE_NAME));
151+
assertThat(SESSION_RULE.session().getMetadata().getTabletMap().getMapping().keySet())
152+
.doesNotContain(TABLET_MAP_KEY);
153+
}
154+
155+
@Test
156+
public void should_remove_tablets_on_table_update() {
157+
SESSION_RULE
158+
.session()
159+
.execute("ALTER TABLE " + KEYSPACE_NAME + "." + TABLE_NAME + " ADD anotherCol int");
160+
ArgumentCaptor<TableMetadata> previous = ArgumentCaptor.forClass(TableMetadata.class);
161+
Mockito.verify(listener, Mockito.timeout(NOTIF_TIMEOUT_MS).times(1))
162+
.onTableUpdated(Mockito.any(), previous.capture());
163+
assertThat(previous.getValue().getName()).isEqualTo(CqlIdentifier.fromCql(TABLE_NAME));
164+
assertThat(SESSION_RULE.session().getMetadata().getTabletMap().getMapping().keySet())
165+
.doesNotContain(TABLET_MAP_KEY);
166+
}
167+
168+
@Test
169+
public void should_remove_tablets_on_table_drop() {
170+
SESSION_RULE.session().execute("DROP TABLE " + KEYSPACE_NAME + "." + TABLE_NAME);
171+
ArgumentCaptor<TableMetadata> table = ArgumentCaptor.forClass(TableMetadata.class);
172+
Mockito.verify(listener, Mockito.timeout(NOTIF_TIMEOUT_MS).times(1))
173+
.onTableDropped(table.capture());
174+
assertThat(table.getValue().getName()).isEqualTo(CqlIdentifier.fromCql(TABLE_NAME));
175+
assertThat(SESSION_RULE.session().getMetadata().getTabletMap().getMapping().keySet())
176+
.doesNotContain(TABLET_MAP_KEY);
177+
}
178+
}

0 commit comments

Comments
 (0)