Skip to content

Commit d679372

Browse files
rvansarvansa
committed
Support Session suspend and resume
If an application wants to use OpenJDK CRaC it must terminate all connections to nodes before checkpoint. Here we expose a high-level API in SessionLifecycleManager without relying on CRaC itself.
1 parent f42ab99 commit d679372

File tree

3 files changed

+260
-0
lines changed

3 files changed

+260
-0
lines changed

‎core/src/main/java/com/datastax/oss/driver/api/core/session/SessionLifecycleManager.java

+69
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
package com.datastax.oss.driver.api.core.session;
19+
20+
import com.datastax.oss.driver.internal.core.session.DefaultSession;
21+
import com.datastax.oss.driver.internal.core.session.DefaultSessionLifecycleManager;
22+
import java.util.concurrent.CompletionStage;
23+
24+
/**
25+
* Provides extra methods for {@link Session} lifecycle. In a suspended state the session should not
26+
* keep any connections to cluster nodes open.
27+
*/
28+
public interface SessionLifecycleManager {
29+
/**
30+
* Creates a new manager for session lifecycle.
31+
*
32+
* @param session Session that should be managed.
33+
* @return The new session lifecycle manager.
34+
* @throws IllegalArgumentException if the session cannot be managed.
35+
*/
36+
static SessionLifecycleManager of(Session session) {
37+
if (session instanceof DefaultSession) {
38+
return new DefaultSessionLifecycleManager((DefaultSession) session);
39+
} else {
40+
throw new IllegalArgumentException(session + " is not an instance of DefaultSession");
41+
}
42+
}
43+
44+
/**
45+
* Terminates all connections to cluster nodes.
46+
*
47+
* @return Stage that completes when all connections are terminated.
48+
*/
49+
CompletionStage<Void> suspendAsync();
50+
51+
/** Helper method invoking {@link #suspendAsync()} in a synchronous way. */
52+
default void suspend() {
53+
suspendAsync().toCompletableFuture().join();
54+
}
55+
56+
/**
57+
* Triggers reconnection to the cluster. This reconnection proceeds asynchronously; the invocation
58+
* does not wait for connections establishment.
59+
*
60+
* @return
61+
*/
62+
void resume();
63+
64+
/**
65+
* @return True if the session is {@link #suspendAsync()} was called, until {@link #resumeAsync()}
66+
* is called.
67+
*/
68+
boolean isSuspended();
69+
}

‎core/src/main/java/com/datastax/oss/driver/internal/core/session/DefaultSessionLifecycleManager.java

+119
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
package com.datastax.oss.driver.internal.core.session;
19+
20+
import com.datastax.oss.driver.api.core.metadata.Node;
21+
import com.datastax.oss.driver.api.core.metadata.NodeState;
22+
import com.datastax.oss.driver.api.core.session.SessionLifecycleManager;
23+
import com.datastax.oss.driver.internal.core.context.InternalDriverContext;
24+
import com.datastax.oss.driver.internal.core.metadata.DefaultNode;
25+
import com.datastax.oss.driver.internal.core.metadata.NodeStateEvent;
26+
import com.datastax.oss.driver.internal.core.pool.ChannelPool;
27+
import com.datastax.oss.driver.internal.core.util.concurrent.CompletableFutures;
28+
import java.util.ArrayList;
29+
import java.util.HashMap;
30+
import java.util.Map;
31+
import java.util.concurrent.CompletableFuture;
32+
import java.util.concurrent.CompletionStage;
33+
34+
public class DefaultSessionLifecycleManager implements SessionLifecycleManager {
35+
private final DefaultSession session;
36+
private final InternalDriverContext context;
37+
private CompletableFuture<Void> suspendFuture;
38+
private Map<Node, NodeState> lastState;
39+
40+
public DefaultSessionLifecycleManager(DefaultSession session) {
41+
this.session = session;
42+
this.context = (InternalDriverContext) session.getContext();
43+
}
44+
45+
@Override
46+
public synchronized CompletionStage<Void> suspendAsync() {
47+
if (suspendFuture != null) {
48+
return suspendFuture;
49+
}
50+
suspendFuture = new CompletableFuture<>();
51+
// ControlConnection would try to reconnect when it receives the event that
52+
// node was brought down; closing the very channel to this node prevents that.
53+
this.context
54+
.getControlConnection()
55+
.channel()
56+
.close()
57+
.addListener(
58+
f -> {
59+
if (f.isSuccess()) {
60+
forceNodesDown()
61+
.whenComplete(
62+
(ignored, throwable) -> {
63+
if (throwable != null) {
64+
suspendFuture.completeExceptionally(throwable);
65+
} else {
66+
suspendFuture.complete(null);
67+
}
68+
});
69+
} else {
70+
suspendFuture.completeExceptionally(f.cause());
71+
}
72+
});
73+
return suspendFuture;
74+
}
75+
76+
private CompletionStage<Void> forceNodesDown() {
77+
lastState = new HashMap<>();
78+
ArrayList<CompletionStage<Void>> closeFutures = new ArrayList<>();
79+
for (Map.Entry<Node, ChannelPool> e : session.getPools().entrySet()) {
80+
Node node = e.getKey();
81+
NodeState currentState = node.getState();
82+
lastState.put(node, currentState);
83+
closeFutures.add(e.getValue().closeFuture());
84+
context
85+
.getEventBus()
86+
.fire(NodeStateEvent.changed(currentState, NodeState.FORCED_DOWN, (DefaultNode) node));
87+
}
88+
return CompletableFutures.allDone(closeFutures);
89+
}
90+
91+
@Override
92+
public void resume() {
93+
if (suspendFuture == null) {
94+
return;
95+
}
96+
suspendFuture.whenComplete(
97+
(ignored, throwable) -> {
98+
if (throwable != null || lastState == null) {
99+
return;
100+
}
101+
synchronized (this) {
102+
for (Map.Entry<Node, NodeState> e : lastState.entrySet()) {
103+
NodeStateEvent changed =
104+
NodeStateEvent.changed(
105+
NodeState.FORCED_DOWN, e.getValue(), (DefaultNode) e.getKey());
106+
this.context.getEventBus().fire(changed);
107+
}
108+
lastState = null;
109+
suspendFuture = null;
110+
context.getControlConnection().reconnectNow();
111+
}
112+
});
113+
}
114+
115+
@Override
116+
public synchronized boolean isSuspended() {
117+
return suspendFuture != null;
118+
}
119+
}

‎integration-tests/src/test/java/com/datastax/oss/driver/core/session/SuspendIT.java

+72
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
package com.datastax.oss.driver.core.session;
19+
20+
import static com.datastax.oss.simulacron.common.stubbing.PrimeDsl.noRows;
21+
import static com.datastax.oss.simulacron.common.stubbing.PrimeDsl.when;
22+
import static org.assertj.core.api.Assertions.assertThat;
23+
import static org.assertj.core.api.Assertions.assertThatThrownBy;
24+
25+
import com.datastax.oss.driver.api.core.CqlSession;
26+
import com.datastax.oss.driver.api.core.NoNodeAvailableException;
27+
import com.datastax.oss.driver.api.core.session.SessionLifecycleManager;
28+
import com.datastax.oss.driver.api.testinfra.session.SessionUtils;
29+
import com.datastax.oss.driver.api.testinfra.simulacron.SimulacronRule;
30+
import com.datastax.oss.driver.categories.ParallelizableTests;
31+
import com.datastax.oss.driver.internal.core.context.InternalDriverContext;
32+
import com.datastax.oss.driver.internal.core.session.PoolManager;
33+
import com.datastax.oss.simulacron.common.cluster.ClusterSpec;
34+
import org.junit.ClassRule;
35+
import org.junit.Test;
36+
import org.junit.experimental.categories.Category;
37+
38+
@Category(ParallelizableTests.class)
39+
public class SuspendIT {
40+
@ClassRule
41+
public static final SimulacronRule SIMULACRON_RULE =
42+
new SimulacronRule(ClusterSpec.builder().withNodes(2));
43+
44+
private static final String QUERY_STRING = "select * from foo";
45+
46+
@Test
47+
public void should_resume_after_suspend() throws Exception {
48+
SIMULACRON_RULE.cluster().prime(when(QUERY_STRING).then(noRows()));
49+
50+
CqlSession session = SessionUtils.newSession(SIMULACRON_RULE);
51+
assertThat(session.execute(QUERY_STRING).all().size()).isEqualTo(0);
52+
53+
SessionLifecycleManager manager = SessionLifecycleManager.of(session);
54+
manager.suspend();
55+
56+
PoolManager poolManager = ((InternalDriverContext) session.getContext()).getPoolManager();
57+
assertThat(poolManager.getPools().size()).isEqualTo(0);
58+
assertThatThrownBy(() -> session.execute(QUERY_STRING).all())
59+
.isInstanceOf(NoNodeAvailableException.class);
60+
61+
manager.resume();
62+
63+
// Busy waiting - PoolManager does not expose any listeners on added node.
64+
// After ChannelEvent.Type.OPEN the future is added to PoolManager.SingleThreaded.pending
65+
// but this map is not exposed (and not synchronized)
66+
while (poolManager.getPools().size() == 0) {
67+
Thread.sleep(10);
68+
}
69+
70+
assertThat(session.execute(QUERY_STRING).all().size()).isEqualTo(0);
71+
}
72+
}

0 commit comments

Comments
 (0)