/* * Copyright (C) 2017-2017 DataStax Inc. * * Licensed 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 com.datastax.oss.simulacron.server; import static com.datastax.oss.protocol.internal.response.result.Void.INSTANCE; import static com.datastax.oss.simulacron.common.stubbing.DisconnectAction.Scope.CLUSTER; import static com.datastax.oss.simulacron.common.stubbing.DisconnectAction.Scope.NODE; import static com.datastax.oss.simulacron.common.stubbing.PrimeDsl.noRows; import static com.datastax.oss.simulacron.common.stubbing.PrimeDsl.when; import static com.datastax.oss.simulacron.common.utils.FrameUtils.wrapResponse; import static com.datastax.oss.simulacron.server.ChannelUtils.completable; import static com.datastax.oss.simulacron.server.FrameCodecUtils.buildFrameCodec; import com.datastax.oss.protocol.internal.Frame; import com.datastax.oss.protocol.internal.Message; import com.datastax.oss.protocol.internal.request.Batch; import com.datastax.oss.protocol.internal.request.Execute; import com.datastax.oss.protocol.internal.request.Options; import com.datastax.oss.protocol.internal.request.Prepare; import com.datastax.oss.protocol.internal.request.Query; import com.datastax.oss.protocol.internal.request.Register; import com.datastax.oss.protocol.internal.request.Startup; import com.datastax.oss.protocol.internal.response.Ready; import com.datastax.oss.protocol.internal.response.Supported; import com.datastax.oss.protocol.internal.response.error.Unprepared; import com.datastax.oss.protocol.internal.response.result.SetKeyspace; import com.datastax.oss.simulacron.common.cluster.AbstractNode; import com.datastax.oss.simulacron.common.cluster.ActivityLog; import com.datastax.oss.simulacron.common.cluster.ClusterConnectionReport; import com.datastax.oss.simulacron.common.cluster.ClusterQueryLogReport; import com.datastax.oss.simulacron.common.cluster.NodeConnectionReport; import com.datastax.oss.simulacron.common.cluster.NodeQueryLogReport; import com.datastax.oss.simulacron.common.cluster.NodeSpec; import com.datastax.oss.simulacron.common.cluster.QueryLog; import com.datastax.oss.simulacron.common.stubbing.Action; import com.datastax.oss.simulacron.common.stubbing.CloseType; import com.datastax.oss.simulacron.common.stubbing.DisconnectAction; import com.datastax.oss.simulacron.common.stubbing.MessageResponseAction; import com.datastax.oss.simulacron.common.stubbing.NoResponseAction; import com.datastax.oss.simulacron.common.stubbing.Prime; import com.datastax.oss.simulacron.common.stubbing.PrimeDsl.PrimeBuilder; import com.datastax.oss.simulacron.common.stubbing.StubMapping; import com.datastax.oss.simulacron.server.listener.QueryListener; import com.fasterxml.jackson.annotation.JsonIgnore; import io.netty.bootstrap.ServerBootstrap; import io.netty.channel.Channel; import io.netty.channel.ChannelFuture; import io.netty.channel.ChannelFutureListener; import io.netty.channel.ChannelHandlerContext; import io.netty.channel.group.ChannelGroup; import io.netty.channel.group.DefaultChannelGroup; import io.netty.channel.socket.SocketChannel; import io.netty.util.Timeout; import io.netty.util.Timer; import io.netty.util.TimerTask; import io.netty.util.concurrent.GlobalEventExecutor; import java.math.BigInteger; import java.net.SocketAddress; import java.util.ArrayList; import java.util.Arrays; import java.util.Collection; import java.util.Collections; import java.util.HashMap; import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.Optional; import java.util.UUID; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionStage; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Function; import java.util.function.Predicate; import java.util.regex.Matcher; import java.util.regex.Pattern; import java.util.stream.Collectors; import java.util.stream.Stream; import org.slf4j.Logger; import org.slf4j.LoggerFactory; public class BoundNode extends AbstractNode<BoundCluster, BoundDataCenter> implements BoundTopic<NodeConnectionReport, NodeQueryLogReport> { static final Predicate<QueryLog> ALWAYS_TRUE = x -> true; private static Logger logger = LoggerFactory.getLogger(BoundNode.class); private static final Pattern useKeyspacePattern = Pattern.compile("\\s*use\\s+(.*)$", Pattern.CASE_INSENSITIVE); private final transient ServerBootstrap bootstrap; // TODO: Isn't really a good reason for this to be an AtomicReference as if binding fails we don't // reset // the channel, but leaving it this way for now in case there is a future use case. final transient AtomicReference<Channel> channel; final transient ChannelGroup clientChannelGroup = new DefaultChannelGroup(GlobalEventExecutor.INSTANCE); // TODO: There could be a lot of concurrency issues around simultaneous calls to reject/accept, // however in the general case we don't expect it. Leave this as AtomicReference in case we want // to handle it better. private final transient AtomicReference<RejectState> rejectState = new AtomicReference<>(new RejectState()); private final transient Timer timer; private final transient StubStore stubStore; private final boolean activityLogging; private final Server server; private final BoundCluster cluster; private final transient List<QueryListenerWrapper> queryListeners = new ArrayList<>(); final transient ActivityLog activityLog = new ActivityLog(); private final transient FrameCodecWrapper frameCodec; private static class RejectState { private final RejectScope scope; private volatile int rejectAfter; private volatile boolean listeningForNewConnections; RejectState() { this(true, Integer.MIN_VALUE, null); } RejectState(boolean listeningForNewConnections, int rejectAfter, RejectScope scope) { this.listeningForNewConnections = listeningForNewConnections; this.rejectAfter = rejectAfter; this.scope = scope; } } BoundNode( SocketAddress address, NodeSpec delegate, Map<String, Object> peerInfo, BoundCluster cluster, BoundDataCenter parent, Server server, Timer timer, Channel channel, boolean activityLogging) { super( address, delegate.getName(), delegate.getId() != null ? delegate.getId() : 0, delegate.getHostId() != null ? delegate.getHostId() : UUID.randomUUID(), delegate.getCassandraVersion(), delegate.getDSEVersion(), peerInfo, parent); this.cluster = cluster; this.server = server; // for test purposes server may be null. this.bootstrap = server != null ? server.serverBootstrap : null; this.timer = timer; this.channel = new AtomicReference<>(channel); this.stubStore = new StubStore(); this.activityLogging = activityLogging; this.frameCodec = buildFrameCodec(delegate).orElse(parent.getFrameCodec()); } @Override public Long getActiveConnections() { // Filter only active channels as some may be in process of closing. return clientChannelGroup.stream().filter(Channel::isActive).count(); } /** * Closes the listening channel for this node. Note that this does not close existing client * connections, this can be done using {@link #disconnectConnections()}. To stop listening and * close connections, use {@link #close()}. * * @return future that completes when listening channel is closed. */ private CompletableFuture<Void> unbind() { logger.debug("Unbinding listener on {}", channel); return completable(channel.get().close()).thenApply(v -> null); } /** * Reopens the listening channel for this node. If the channel was already open, has no effect and * future completes immediately. * * @return future that completes when listening channel is reopened. */ private CompletableFuture<Void> rebind() { if (this.channel.get().isOpen()) { // already accepting... return CompletableFuture.completedFuture(null); } CompletableFuture<Void> future = new CompletableFuture<>(); ChannelFuture bindFuture = bootstrap.bind(this.getAddress()); bindFuture.addListener( (ChannelFutureListener) channelFuture -> { if (channelFuture.isSuccess()) { channelFuture.channel().attr(Server.HANDLER).set(this); logger.debug("Bound {} to {}", BoundNode.this, channelFuture.channel()); future.complete(null); channel.set(channelFuture.channel()); } else { // If failed, propagate it. future.completeExceptionally( new BindNodeException(BoundNode.this, getAddress(), channelFuture.cause())); } }); return future; } /** * Disconnects all client channels. Does not close listening interface (see {@link #unbind()} for * that). * * @return future that completes when all client channels are disconnected. */ private CompletionStage<Void> disconnectConnections() { return completable(clientChannelGroup.disconnect()).thenApply(v -> null); } /** * Indicates that the node should resume accepting connections. * * @return future that completes when node is listening again. */ @Override public CompletionStage<Void> acceptConnectionsAsync() { logger.debug("Accepting New Connections"); rejectState.set(new RejectState()); // Reopen listening interface if not currently open. if (!channel.get().isOpen()) { return rebind(); } else { return CompletableFuture.completedFuture(null); } } /** * This is used to fetch the QueryLogReport with unfiltered activity logs * * @return QueryLogReport containing all the logs for the node */ @Override @JsonIgnore public NodeQueryLogReport getLogs() { ClusterQueryLogReport clusterQueryLogReportReport = new ClusterQueryLogReport(cluster.getId()); return clusterQueryLogReportReport.addNode(this, this.activityLog.getLogs()); } /** * This is used to fetch the QueryLogReport with filtered activity logs * * @return QueryLogReport containing all the logs for the node */ @Override @JsonIgnore public NodeQueryLogReport getLogs(boolean primed) { ClusterQueryLogReport clusterQueryLogReportReport = new ClusterQueryLogReport(cluster.getId()); return clusterQueryLogReportReport.addNode(this, this.activityLog.getLogs(primed)); } @Override public void clearLogs() { activityLog.clear(); } @Override public void registerQueryListener( QueryListener queryListener, boolean after, Predicate<QueryLog> filter) { queryListeners.add(new QueryListenerWrapper(queryListener, after, filter)); } /** * Indicates that the node should stop accepting new connections. * * @param after If non-zero, after how many successful startup messages should stop accepting * connections. * @param scope The scope to reject connections, either stop listening for connections, or accept * connections but don't respond to startup requests. * @return future that completes when listening channel is unbound (if {@link RejectScope#UNBIND} * was used) or immediately if {@link RejectScope#REJECT_STARTUP} was used or after > 0. */ @Override public CompletionStage<Void> rejectConnectionsAsync(int after, RejectScope scope) { RejectState state; if (after <= 0) { logger.debug("Rejecting new connections with scope {}", scope); state = new RejectState(false, Integer.MIN_VALUE, scope); } else { logger.debug("Rejecting new connections after {} attempts with scope {}", after, scope); state = new RejectState(true, after, scope); } rejectState.set(state); if (after <= 0 && scope != RejectScope.REJECT_STARTUP) { CompletableFuture<Void> unbindFuture = unbind(); // if scope is STOP, disconnect existing connections after unbinding. if (scope == RejectScope.STOP) { return unbindFuture.thenCompose(n -> disconnectConnections()); } else { return unbindFuture; } } else { return CompletableFuture.completedFuture(null); } } /** * Search stub stores for matches for the given frame and this node. If not found at node level, * checks dc level, if not found at dc level, checks cluster level, if not found at cluster level, * checks global store. * * @param frame frame to match on. * @return matching stub if present. */ private Optional<StubMapping> find(Frame frame) { Optional<StubMapping> stub = stubStore.find(this, frame); if (!stub.isPresent()) { return getDataCenter().find(this, frame); } return stub; } void handle(ChannelHandlerContext ctx, UnsupportedProtocolVersionMessage message) { if (activityLogging) { QueryLog queryLog = activityLog.addLog( message.getFrame(), ctx.channel().remoteAddress(), System.currentTimeMillis(), Optional.empty()); notifyQueryListeners(queryLog, false); } } void handle(ChannelHandlerContext ctx, Frame frame) { logger.debug("Got request streamId: {} msg: {}", frame.streamId, frame.message); // On receiving a message, first check the stub store to see if there is handling logic for it. // If there is, handle each action. // Otherwise delegate to default behavior. Optional<StubMapping> stubOption = find(frame); List<Action> actions = null; if (stubOption.isPresent()) { StubMapping stub = stubOption.get(); actions = stub.getActions(this, frame); } QueryLog queryLog = null; // store the frame in history if (activityLogging) { queryLog = activityLog.addLog( frame, ctx.channel().remoteAddress(), System.currentTimeMillis(), stubOption); notifyQueryListeners(queryLog, false); } if (actions != null && !actions.isEmpty()) { // TODO: It might be useful to tie behavior to completion of actions but for now this isn't // necessary. CompletableFuture<Void> future = new CompletableFuture<>(); handleActions(actions.iterator(), ctx, frame, future, queryLog); } else { // Future that if set defers sending the message until the future completes. CompletableFuture<?> deferFuture = null; Message response = null; if (frame.message instanceof Startup || frame.message instanceof Register) { RejectState state = rejectState.get(); // We aren't listening for new connections, return immediately. if (!state.listeningForNewConnections) { return; } else if (state.rejectAfter > 0) { // Decrement rejectAfter indicating a new initialization attempt. state.rejectAfter--; if (state.rejectAfter == 0) { // If reject after is now 0, indicate that it's time to stop listening (but allow this // one) state.rejectAfter = -1; state.listeningForNewConnections = false; deferFuture = rejectConnectionsAsync(-1, state.scope).toCompletableFuture(); } } response = new Ready(); } else if (frame.message instanceof Options) { // Maybe eventually we can set these depending on the version but so far it looks // like this.cassandraVersion and this.dseVersion are both null HashMap<String, List<String>> options = new HashMap<>(); options.put("PROTOCOL_VERSIONS", Arrays.asList("3/v3", "4/v4", "5/v5-beta")); options.put("CQL_VERSION", Collections.singletonList("3.4.4")); options.put("COMPRESSION", Arrays.asList("snappy", "lz4")); response = new Supported(options); } else if (frame.message instanceof Query) { Query query = (Query) frame.message; String queryStr = query.query; if (queryStr.startsWith("USE") || queryStr.startsWith("use")) { Matcher matcher = useKeyspacePattern.matcher(queryStr); // should always match. assert matcher.matches(); if (matcher.matches()) { // unquote keyspace if quoted, cassandra doesn't expect keyspace to be quoted coming // back String keyspace = matcher.group(1).replaceAll("^\"|\"$", ""); response = new SetKeyspace(keyspace); } } else { response = INSTANCE; } } else if (frame.message instanceof Batch) { response = INSTANCE; } else if (frame.message instanceof Execute) { // Unprepared execute received, return an unprepared. Execute execute = (Execute) frame.message; String hex = new BigInteger(1, execute.queryId).toString(16); response = new Unprepared("No prepared statement with id: " + hex, execute.queryId); } else if (frame.message instanceof Prepare) { // fake up a prepared statement from the message and register an internal prime for it. Prepare prepare = (Prepare) frame.message; // TODO: Maybe attempt to identify bind parameters String query = prepare.cqlQuery; Prime prime = whenWithInferredParams(query).then(noRows()).build(); this.getCluster().getStubStore().registerInternal(prime); response = prime.toPrepared(); } if (response != null) { final QueryLog fQueryLog = queryLog; if (deferFuture != null) { final Message fResponse = response; deferFuture.thenRun( () -> { sendMessage(ctx, frame, fResponse) .addListener( (x) -> { notifyQueryListeners(fQueryLog, true); }); }); } else { sendMessage(ctx, frame, response) .addListener((x) -> notifyQueryListeners(fQueryLog, true)); } } else { notifyQueryListeners(queryLog, true); } } } private void notifyQueryListeners(QueryLog queryLog, boolean after) { if (queryLog != null && !queryListeners.isEmpty()) { for (QueryListenerWrapper wrapper : queryListeners) { if (after == wrapper.after) { wrapper.apply(this, queryLog); } } } getDataCenter().notifyQueryListeners(this, queryLog, after); } private void handleActions( Iterator<Action> nextActions, ChannelHandlerContext ctx, Frame frame, CompletableFuture<Void> doneFuture, QueryLog queryLog) { // If there are no more actions, complete the done future and return. if (!nextActions.hasNext()) { doneFuture.complete(null); notifyQueryListeners(queryLog, true); return; } CompletableFuture<Void> future = new CompletableFuture<>(); Action action = nextActions.next(); ActionHandler handler = new ActionHandler(action, ctx, frame, future); if (action.delayInMs() > 0) { timer.newTimeout(handler, action.delayInMs(), TimeUnit.MILLISECONDS); } else { // process immediately when delay is 0. handler.run(null); } // proceed to next action when complete future.whenComplete( (v, ex) -> { if (ex != null) { doneFuture.completeExceptionally(ex); } else { handleActions(nextActions, ctx, frame, doneFuture, queryLog); } }); } private class ActionHandler implements TimerTask { private final Action action; private final ChannelHandlerContext ctx; private final Frame frame; private final CompletableFuture<Void> doneFuture; ActionHandler( Action action, ChannelHandlerContext ctx, Frame frame, CompletableFuture<Void> doneFuture) { this.action = action; this.ctx = ctx; this.frame = frame; this.doneFuture = doneFuture; } @Override public void run(Timeout timeout) { CompletableFuture<Void> future; // TODO maybe delegate this logic elsewhere if (action instanceof MessageResponseAction) { MessageResponseAction mAction = (MessageResponseAction) action; future = completable(sendMessage(ctx, frame, mAction.getMessage())); } else if (action instanceof DisconnectAction) { DisconnectAction cAction = (DisconnectAction) action; switch (cAction.getScope()) { case CONNECTION: future = closeConnectionAsync(ctx.channel().remoteAddress(), cAction.getCloseType()) .toCompletableFuture() .thenApply(v -> null); break; default: Stream<BoundNode> nodes = cAction.getScope() == NODE ? Stream.of(BoundNode.this) : cAction.getScope() == CLUSTER ? getCluster().getNodes().stream() : getDataCenter().getNodes().stream(); future = closeNodes(nodes, cAction.getCloseType()); break; } } else if (action instanceof NoResponseAction) { future = new CompletableFuture<>(); future.complete(null); } else { logger.warn("Got action {} that we don't know how to handle.", action); future = new CompletableFuture<>(); future.complete(null); } future.whenComplete( (v, t) -> { if (t != null) { doneFuture.completeExceptionally(t); } else { doneFuture.complete(v); } }); } } private static CompletableFuture<Void> closeNodes(Stream<BoundNode> nodes, CloseType closeType) { return CompletableFuture.allOf( nodes .map(n -> n.closeConnectionsAsync(closeType).toCompletableFuture()) .collect(Collectors.toList()) .toArray(new CompletableFuture[] {})); } private ChannelFuture sendMessage( ChannelHandlerContext ctx, Frame requestFrame, Message responseMessage) { Frame responseFrame = wrapResponse(requestFrame, responseMessage); logger.debug( "Sending response for streamId: {} with msg {}", responseFrame.streamId, responseFrame.message); return ctx.writeAndFlush(responseFrame); } @Override public StubStore getStubStore() { return stubStore; } @Override public int clearPrimes(boolean nested) { return stubStore.clear(); } @Override public CompletionStage<BoundCluster> unregisterAsync() { return getServer().unregisterAsync(this); } /** See {@link #clearPrimes(boolean)} */ public int clearPrimes() { return stubStore.clear(); } @Override public NodeConnectionReport getConnections() { ClusterConnectionReport clusterConnectionReport = new ClusterConnectionReport(cluster.getId()); return clusterConnectionReport.addNode( this, clientChannelGroup.stream().map(Channel::remoteAddress).collect(Collectors.toList()), getAddress()); } @Override public CompletionStage<NodeConnectionReport> closeConnectionsAsync(CloseType closeType) { NodeConnectionReport report = getConnections(); return closeChannelGroup(this.clientChannelGroup, closeType).thenApply(v -> report); } @Override public NodeConnectionReport pauseRead() { this.clientChannelGroup.forEach(c -> c.config().setAutoRead(false)); return getConnections(); } @Override public NodeConnectionReport resumeRead() { this.clientChannelGroup.forEach(c -> c.config().setAutoRead(true)); return getConnections(); } private static CompletableFuture<Void> closeChannelGroup( ChannelGroup channelGroup, CloseType closeType) { switch (closeType) { case DISCONNECT: return completable(channelGroup.disconnect()); default: return CompletableFuture.allOf( channelGroup .stream() .map( c -> { CompletableFuture<Void> f; Function<SocketChannel, ChannelFuture> shutdownMethod = closeType == CloseType.SHUTDOWN_READ ? SocketChannel::shutdownInput : SocketChannel::shutdownOutput; if (c instanceof SocketChannel) { f = completable(shutdownMethod.apply((SocketChannel) c)); } else { logger.warn( "Got {} request for non-SocketChannel {}, disconnecting instead.", closeType, c); f = completable(c.disconnect()); } return f; }) .collect(Collectors.toList()) .toArray(new CompletableFuture[] {})); } } @Override public CompletionStage<NodeConnectionReport> closeConnectionAsync( SocketAddress connection, CloseType type) { Optional<Channel> channel = this.clientChannelGroup .stream() .filter(c -> c.remoteAddress().equals(connection)) .findFirst(); if (channel.isPresent()) { ChannelGroup channelGroup = new DefaultChannelGroup(GlobalEventExecutor.INSTANCE); channelGroup.add(channel.get()); ClusterConnectionReport clusterReport = new ClusterConnectionReport(getCluster().getId()); NodeConnectionReport report = clusterReport.addNode(this, Collections.singletonList(connection), getAddress()); return closeChannelGroup(channelGroup, type).thenApply(f -> report); } else { CompletableFuture<NodeConnectionReport> failedFuture = new CompletableFuture<>(); failedFuture.completeExceptionally(new IllegalArgumentException("Not found")); return failedFuture; } } @Override public Collection<BoundNode> getNodes() { return Collections.singleton(this); } @Override public Server getServer() { return server; } @Override @JsonIgnore public FrameCodecWrapper getFrameCodec() { return frameCodec; } /** * Convenience fluent builder for constructing a prime with a query, where the parameters are * inferred by the query * * @param query The query string to match against. * @return builder for this prime. */ private static PrimeBuilder whenWithInferredParams(String query) { long posParamCount = query.chars().filter(num -> num == '?').count(); // Do basic param population for positional types HashMap<String, String> paramTypes = new HashMap<>(); HashMap<String, Object> params = new HashMap<>(); if (posParamCount > 0) { for (int i = 0; i < posParamCount; i++) { params.put(Integer.toString(i), "*"); paramTypes.put(Integer.toString(i), "varchar"); } } // Do basic param population for named types else { List<String> allMatches = new ArrayList<>(); Pattern p = Pattern.compile("([\\w']+)\\s=\\s:[\\w]+"); Matcher m = p.matcher(query); while (m.find()) { allMatches.add(m.group(1)); } for (String match : allMatches) { params.put(match, "*"); paramTypes.put(match, "varchar"); } } return when( new com.datastax.oss.simulacron.common.request.Query( query, Collections.emptyList(), params, paramTypes)); } }