Skip to content

Commit 95fa50b

Browse files
BouncheckBouncheck
committed
Remove tablet mappings on select events
Registers an event listener dedicated to updating the `TabletMap` held by `Metadata`. It will trigger removal of relevant mappings whenever it's notified of any change to the table schema or keyspace schema including dropping them. Introduces an integration test that verifies this listener's behaviour. Addresses #378.
1 parent 1cfc2df commit 95fa50b

File tree

4 files changed

+271
-0
lines changed

4 files changed

+271
-0
lines changed

‎driver-core/src/main/java/com/datastax/driver/core/Metadata.java

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -150,6 +150,22 @@ public class Metadata {
150150
Metadata(Cluster.Manager cluster) {
151151
this.cluster = cluster;
152152
this.tabletMap = TabletMap.emptyMap(cluster);
153+
if (cluster.configuration.getQueryOptions().isMetadataEnabled()) {
154+
cluster
155+
.getCluster()
156+
.register(
157+
new TabletMapListener(tabletMap) {
158+
@Override
159+
public void onRegister(Cluster cluster) {
160+
logger.info("Registered event listener for tablet map.");
161+
}
162+
163+
@Override
164+
public void onUnregister(Cluster cluster) {
165+
logger.info("Unregistered tablet map's event listener.");
166+
}
167+
});
168+
}
153169
}
154170

155171
// rebuilds the token map with the current hosts, typically when refreshing schema metadata

‎driver-core/src/main/java/com/datastax/driver/core/TabletMap.java

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -198,6 +198,24 @@ public TypeCodec<TupleValue> getTabletPayloadCodec() {
198198
return tabletPayloadCodec;
199199
}
200200

201+
public void removeTableMappings(KeyspaceTableNamePair key) {
202+
this.mapping.remove(key);
203+
}
204+
205+
public void removeTableMappings(String keyspace, String table) {
206+
removeTableMappings(new KeyspaceTableNamePair(keyspace, table));
207+
}
208+
209+
public void removeTableMappings(String keyspace) {
210+
Iterator<TabletMap.KeyspaceTableNamePair> it = getMapping().keySet().iterator();
211+
while (it.hasNext()) {
212+
KeyspaceTableNamePair key = it.next();
213+
if (key.getKeyspace().equals(keyspace)) {
214+
it.remove();
215+
}
216+
}
217+
}
218+
201219
/**
202220
* Simple class to hold UUID of a host and a shard number. Class itself makes no checks or
203221
* guarantees about existence of a shard on corresponding host.

‎driver-core/src/main/java/com/datastax/driver/core/TabletMapListener.java

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
package com.datastax.driver.core;
2+
3+
public class TabletMapListener extends SchemaChangeListenerBase {
4+
private final TabletMap tabletMap;
5+
6+
public TabletMapListener(TabletMap tabletMap) {
7+
this.tabletMap = tabletMap;
8+
}
9+
10+
@Override
11+
public void onTableChanged(TableMetadata current, TableMetadata previous) {
12+
tabletMap.removeTableMappings(previous.getKeyspace().getName(), previous.getName());
13+
}
14+
15+
@Override
16+
public void onTableRemoved(TableMetadata table) {
17+
tabletMap.removeTableMappings(table.getKeyspace().getName(), table.getName());
18+
}
19+
20+
@Override
21+
public void onKeyspaceRemoved(KeyspaceMetadata keyspace) {
22+
tabletMap.removeTableMappings(keyspace.getName());
23+
}
24+
25+
@Override
26+
public void onKeyspaceChanged(KeyspaceMetadata current, KeyspaceMetadata previous) {
27+
tabletMap.removeTableMappings(previous.getName());
28+
}
29+
}

‎driver-core/src/test/java/com/datastax/driver/core/TabletMapListenerTest.java

Lines changed: 208 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,208 @@
1+
package com.datastax.driver.core;
2+
3+
import static com.datastax.driver.core.Assertions.assertThat;
4+
import static com.datastax.driver.core.Metadata.handleId;
5+
import static org.mockito.Matchers.anyObject;
6+
import static org.mockito.Mockito.after;
7+
import static org.mockito.Mockito.mock;
8+
import static org.mockito.Mockito.timeout;
9+
import static org.mockito.Mockito.verify;
10+
11+
import com.datastax.driver.core.policies.RoundRobinPolicy;
12+
import com.datastax.driver.core.utils.ScyllaVersion;
13+
import com.google.common.collect.Lists;
14+
import java.util.List;
15+
import java.util.concurrent.TimeUnit;
16+
import org.mockito.ArgumentCaptor;
17+
import org.testng.annotations.AfterMethod;
18+
import org.testng.annotations.BeforeMethod;
19+
import org.testng.annotations.Test;
20+
21+
@CCMConfig(
22+
createCluster = false,
23+
numberOfNodes = 2,
24+
jvmArgs = {
25+
"--experimental-features=consistent-topology-changes",
26+
"--experimental-features=tablets"
27+
})
28+
@ScyllaVersion(minOSS = "6.0.0", minEnterprise = "2024.2", description = "Needs to support tablets")
29+
public class TabletMapListenerTest extends CCMTestsSupport {
30+
private static final int INITIAL_TABLETS = 32;
31+
private static final String KEYSPACE_NAME = "listenerTest";
32+
private static final String TABLE_NAME = "testTable";
33+
private static final String CREATE_TABLETS_KEYSPACE_QUERY =
34+
"CREATE KEYSPACE "
35+
+ KEYSPACE_NAME
36+
+ " WITH replication = {'class': "
37+
+ "'NetworkTopologyStrategy', "
38+
+ "'replication_factor': '1'} AND durable_writes = true AND tablets = "
39+
+ "{'initial': "
40+
+ INITIAL_TABLETS
41+
+ "};";
42+
43+
private static final String CREATE_KEYSPACE = CREATE_TABLETS_KEYSPACE_QUERY;
44+
private static final String ALTER_KEYSPACE =
45+
"ALTER KEYSPACE " + KEYSPACE_NAME + " WITH durable_writes = false";
46+
private static final String DROP_KEYSPACE = "DROP KEYSPACE " + KEYSPACE_NAME;
47+
48+
private static final String CREATE_TABLE =
49+
"CREATE TABLE " + KEYSPACE_NAME + "." + TABLE_NAME + "(i int primary key)";
50+
private static final String INSERT_QUERY_TEMPLATE =
51+
"INSERT INTO " + KEYSPACE_NAME + "." + TABLE_NAME + " (i) VALUES (%s)";
52+
private static final String INSERT_ALTERED_TEMPLATE =
53+
"INSERT INTO " + KEYSPACE_NAME + "." + TABLE_NAME + " (i,j) VALUES (%s,%s)";
54+
private static final String SELECT_PK_WHERE =
55+
"SELECT i FROM " + KEYSPACE_NAME + "." + TABLE_NAME + " WHERE i = ?";
56+
private static final String ALTER_TABLE =
57+
"ALTER TABLE " + KEYSPACE_NAME + "." + TABLE_NAME + " ADD j int";
58+
private static final String DROP_TABLE = "DROP TABLE " + KEYSPACE_NAME + "." + TABLE_NAME;
59+
60+
/** The maximum time that the test will wait to check that listeners have been notified. */
61+
private static final long NOTIF_TIMEOUT_MS = TimeUnit.MINUTES.toMillis(1);
62+
/** Shorter timeout for less important checks that listeners did not react to specific actions */
63+
private static final long SHORT_TIMEOUT_MS = TimeUnit.SECONDS.toMillis(4);
64+
65+
Cluster cluster;
66+
Session session;
67+
SchemaChangeListener listener;
68+
List<SchemaChangeListener> listeners;
69+
70+
@BeforeMethod(groups = "short")
71+
public void setup() throws InterruptedException {
72+
cluster =
73+
createClusterBuilderNoDebouncing().withLoadBalancingPolicy(new RoundRobinPolicy()).build();
74+
75+
session = cluster.connect();
76+
77+
cluster.register(listener = mock(TabletMapListener.class));
78+
listeners = Lists.newArrayList(listener);
79+
}
80+
81+
// Checks for tablet removal both on table update and removal
82+
@Test(groups = "short")
83+
public void should_remove_tablets_on_table_alterations() throws InterruptedException {
84+
session.execute(CREATE_KEYSPACE);
85+
ArgumentCaptor<KeyspaceMetadata> added = null;
86+
for (SchemaChangeListener listener : listeners) {
87+
added = ArgumentCaptor.forClass(KeyspaceMetadata.class);
88+
verify(listener, timeout(NOTIF_TIMEOUT_MS).times(1)).onKeyspaceAdded(added.capture());
89+
assertThat(added.getValue()).hasName(handleId(KEYSPACE_NAME));
90+
}
91+
assert added != null;
92+
93+
TabletMap tabletMap;
94+
tabletMap = cluster.getMetadata().getTabletMap();
95+
96+
session.execute(CREATE_TABLE);
97+
assertThat(tabletMap.getMapping())
98+
.doesNotContainKey(
99+
new TabletMap.KeyspaceTableNamePair(handleId(KEYSPACE_NAME), handleId(TABLE_NAME)));
100+
101+
session.execute(String.format(INSERT_QUERY_TEMPLATE, "42"));
102+
session.execute(session.prepare(SELECT_PK_WHERE).bind(42));
103+
session.execute(session.prepare(SELECT_PK_WHERE).bind(42));
104+
assertThat(tabletMap.getMapping())
105+
.containsKey(
106+
new TabletMap.KeyspaceTableNamePair(handleId(KEYSPACE_NAME), handleId(TABLE_NAME)));
107+
108+
session.execute(ALTER_TABLE);
109+
for (SchemaChangeListener listener : listeners) {
110+
ArgumentCaptor<TableMetadata> current = ArgumentCaptor.forClass(TableMetadata.class);
111+
ArgumentCaptor<TableMetadata> previous = ArgumentCaptor.forClass(TableMetadata.class);
112+
verify(listener, timeout(NOTIF_TIMEOUT_MS).times(1))
113+
.onTableChanged(current.capture(), previous.capture());
114+
assertThat(previous.getValue().getKeyspace()).hasName(handleId(KEYSPACE_NAME));
115+
assertThat(previous.getValue()).hasName(handleId(TABLE_NAME));
116+
}
117+
assertThat(tabletMap.getMapping())
118+
.doesNotContainKey(
119+
new TabletMap.KeyspaceTableNamePair(handleId(KEYSPACE_NAME), handleId(TABLE_NAME)));
120+
121+
session.execute(String.format(INSERT_ALTERED_TEMPLATE, "42", "42"));
122+
session.execute(session.prepare(SELECT_PK_WHERE).bind(42));
123+
session.execute(session.prepare(SELECT_PK_WHERE).bind(42));
124+
assertThat(tabletMap.getMapping())
125+
.containsKey(
126+
new TabletMap.KeyspaceTableNamePair(handleId(KEYSPACE_NAME), handleId(TABLE_NAME)));
127+
128+
session.execute(DROP_TABLE);
129+
ArgumentCaptor<TableMetadata> removed = null;
130+
for (SchemaChangeListener listener : listeners) {
131+
removed = ArgumentCaptor.forClass(TableMetadata.class);
132+
verify(listener, timeout(NOTIF_TIMEOUT_MS).times(1)).onTableRemoved(removed.capture());
133+
assertThat(removed.getValue().getKeyspace()).hasName(handleId(KEYSPACE_NAME));
134+
assertThat(removed.getValue()).hasName(handleId(TABLE_NAME));
135+
}
136+
assert removed != null;
137+
assertThat(tabletMap.getMapping())
138+
.doesNotContainKey(
139+
new TabletMap.KeyspaceTableNamePair(handleId(KEYSPACE_NAME), handleId(TABLE_NAME)));
140+
141+
session.execute(DROP_KEYSPACE);
142+
}
143+
144+
@Test(groups = "short")
145+
public void should_remove_tablets_on_keyspace_alterations() {
146+
session.execute(CREATE_KEYSPACE);
147+
ArgumentCaptor<KeyspaceMetadata> added = null;
148+
for (SchemaChangeListener listener : listeners) {
149+
added = ArgumentCaptor.forClass(KeyspaceMetadata.class);
150+
verify(listener, timeout(NOTIF_TIMEOUT_MS).times(1)).onKeyspaceAdded(added.capture());
151+
assertThat(added.getValue()).hasName(handleId(KEYSPACE_NAME));
152+
}
153+
assert added != null;
154+
155+
TabletMap tabletMap;
156+
tabletMap = cluster.getMetadata().getTabletMap();
157+
158+
session.execute(CREATE_TABLE);
159+
session.execute(String.format(INSERT_QUERY_TEMPLATE, "42"));
160+
session.execute(session.prepare(SELECT_PK_WHERE).bind(42));
161+
session.execute(session.prepare(SELECT_PK_WHERE).bind(42));
162+
assertThat(tabletMap.getMapping())
163+
.containsKey(
164+
new TabletMap.KeyspaceTableNamePair(handleId(KEYSPACE_NAME), handleId(TABLE_NAME)));
165+
166+
assertThat(cluster.getMetadata().getKeyspace(KEYSPACE_NAME).isDurableWrites()).isTrue();
167+
168+
session.execute(ALTER_KEYSPACE);
169+
assertThat(cluster.getMetadata().getKeyspace(KEYSPACE_NAME)).isNotDurableWrites();
170+
for (SchemaChangeListener listener : listeners) {
171+
ArgumentCaptor<KeyspaceMetadata> current = ArgumentCaptor.forClass(KeyspaceMetadata.class);
172+
ArgumentCaptor<KeyspaceMetadata> previous = ArgumentCaptor.forClass(KeyspaceMetadata.class);
173+
verify(listener, timeout(NOTIF_TIMEOUT_MS).times(1))
174+
.onKeyspaceChanged(current.capture(), previous.capture());
175+
assertThat(previous.getValue()).hasName(handleId(KEYSPACE_NAME));
176+
}
177+
for (SchemaChangeListener listener : listeners) {
178+
verify(listener, after((int) SHORT_TIMEOUT_MS).never())
179+
.onTableChanged(anyObject(), anyObject());
180+
}
181+
assertThat(tabletMap.getMapping())
182+
.doesNotContainKey(
183+
new TabletMap.KeyspaceTableNamePair(handleId(KEYSPACE_NAME), handleId(TABLE_NAME)));
184+
185+
session.execute(session.prepare(SELECT_PK_WHERE).bind(42));
186+
session.execute(session.prepare(SELECT_PK_WHERE).bind(42));
187+
188+
assertThat(tabletMap.getMapping())
189+
.containsKey(
190+
new TabletMap.KeyspaceTableNamePair(handleId(KEYSPACE_NAME), handleId(TABLE_NAME)));
191+
192+
session.execute(DROP_KEYSPACE);
193+
for (SchemaChangeListener listener : listeners) {
194+
ArgumentCaptor<KeyspaceMetadata> removed = ArgumentCaptor.forClass(KeyspaceMetadata.class);
195+
verify(listener, timeout(NOTIF_TIMEOUT_MS).times(1)).onKeyspaceRemoved(removed.capture());
196+
assertThat(removed.getValue()).hasName(handleId(KEYSPACE_NAME));
197+
}
198+
assertThat(tabletMap.getMapping())
199+
.doesNotContainKey(
200+
new TabletMap.KeyspaceTableNamePair(handleId(KEYSPACE_NAME), handleId(TABLE_NAME)));
201+
}
202+
203+
@AfterMethod(groups = "short", alwaysRun = true)
204+
public void teardown() {
205+
if (session != null) session.close();
206+
if (cluster != null) cluster.close();
207+
}
208+
}

0 commit comments

Comments
 (0)