Skip to content

Ensure a Message's Response streamId is always set #4936

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
frankgh wants to merge 7 commits into
base: trunk
Choose a base branch
from
Open

Ensure a Message's Response streamId is always set #4936

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
9b4b177
Ensure a Message's Response streamId is always set
frankgh Jul 13, 2026
7ea6225
Address PR comments from Benedict
frankgh Jul 14, 2026
7614b48
Fix response decoration misses call site from InitialConnectionHandler
frankgh Jul 14, 2026
79e0bf2
Address comments from Benedict
frankgh Jul 15, 2026
8cb7d15
Use inbound.header.streamId
frankgh Jul 15, 2026
ceb5c4e
Cleanup and add javadoc
frankgh Jul 15, 2026
3adeed8
Address issue reported by Caleb
frankgh Jul 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 7 additions & 16 deletions src/java/org/apache/cassandra/transport/CQLMessageHandler.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -465,18 +465,6 @@ private void handleError(Throwable t, int streamId)
errorHandler.accept(ErrorMessage.wrap(t, streamId));
}

/**
* For use in the case where the error can't be mapped to a specific stream id,
* such as a corrupted frame, or when extracting a CQL message from the frame's
* payload fails. This does not attempt to release any resources, as these errors
* should only occur before any capacity acquisition is attempted (e.g. on receipt
* of a corrupt frame, or failure to extract a CQL message from the envelope).
*/
private void handleError(Throwable t)
{
errorHandler.accept(t);
}

// Acts as a Dispatcher.FlushItemConverter
private Framed toFlushItem(Channel channel, Message.Request request, Message.Response response)
{
Expand Down Expand Up @@ -524,8 +512,9 @@ protected boolean processFirstFrameOfLargeMessage(IntactFrame frame, Limit endpo
if (!extracted.isSuccess())
{
// Hard fail on any decoding error as we can't trust the subsequent frames of
// the large message
handleError(ProtocolException.toFatalException(extracted.error()));
// the large message. The stream id is a best-effort value read before extraction
// failed, so route it back where possible rather than defaulting.
handleError(ProtocolException.toFatalException(extracted.error()), extracted.streamId());
return false;
}

Expand All @@ -542,7 +531,7 @@ protected boolean processFirstFrameOfLargeMessage(IntactFrame frame, Limit endpo
// not make sense to continue processing subsequent frames
handleError(ProtocolException.toFatalException(new OversizedAuthMessageException(
MULTI_FRAME_AUTH_ERROR_MESSAGE_PREFIX +
"type = " + header.type + ", size = " + header.bodySizeInBytes)));
"type = " + header.type + ", size = " + header.bodySizeInBytes)), header.streamId);
ClientMetrics.instance.markRequestDiscarded();
return false;
}
Expand Down Expand Up @@ -731,7 +720,9 @@ protected void processCorruptFrame(FrameDecoder.CorruptFrame frame)
processSubsequentFrameOfLargeMessage(frame);
}

handleError(ProtocolException.toFatalException(new ProtocolException(error)));
// A corrupt frame's bytes can't be trusted to map back to a single stream (it may span
// several), so there is no usable stream id here; use the explicit no-request sentinel.
handleError(ProtocolException.toFatalException(new ProtocolException(error)), Message.NO_REQUEST_STREAM_ID);
}

protected void fatalExceptionCaught(Throwable cause)
Expand Down
25 changes: 15 additions & 10 deletions src/java/org/apache/cassandra/transport/Dispatcher.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -112,8 +112,7 @@ public void dispatch(Channel channel, Message.Request request, FlushItemConverte
// We can not respond with a custom, transport, or server exceptions since, given current implementation of clients,
// they will defunct the connection. Without a protocol version bump that introduces an "I am going away message",
// we have to stick to an existing error code.
Message.Response response = ErrorMessage.fromException(new OverloadedException("Server is shutting down"));
response.setStreamId(request.getStreamId());
Message.Response response = ErrorMessage.fromException(new OverloadedException("Server is shutting down"), request.getStreamId());
response.setWarnings(ClientWarn.instance.getWarnings());
response.attach(request.connection);
FlushItem<?> toFlush = forFlusher.toFlushItem(channel, request, response);
Expand Down Expand Up @@ -374,7 +373,7 @@ private static Message.Response processRequest(ServerConnection connection, Mess
if (queueTime > DatabaseDescriptor.getNativeTransportTimeout(TimeUnit.NANOSECONDS))
{
ClientMetrics.instance.markTimedOutBeforeProcessing();
return ErrorMessage.fromException(new OverloadedException("Query timed out before it could start"));
return ErrorMessage.fromException(new OverloadedException("Query timed out before it could start"), request.getStreamId());
}

if (connection.getVersion().isGreaterOrEqualTo(ProtocolVersion.V4))
Expand Down Expand Up @@ -434,8 +433,6 @@ private static Message.Response processRequest(ServerConnection connection, Mess
CoordinatorWriteWarnings.done();
}

response.setStreamId(request.getStreamId());
response.setWarnings(ClientWarn.instance.getWarnings());
response.attach(connection);
connection.applyStateTransition(request.type, response.type);
return response;
Expand All @@ -448,7 +445,10 @@ static Message.Response processRequest(Channel channel, Message.Request request,
{
try
{
return processRequest((ServerConnection) request.connection(), request, backpressure, requestTime);
// decorateResponse must run before the finally below resets ClientWarn, so that it can read the
// warnings accumulated during processing. Decorating in the caller (after this method returns)
// would observe the reset state and drop all client warnings.
return decorateResponse(processRequest((ServerConnection) request.connection(), request, backpressure, requestTime), request);
}
catch (Throwable t)
{
Expand All @@ -461,10 +461,7 @@ static Message.Response processRequest(Channel channel, Message.Request request,
}

Predicate<Throwable> handler = ExceptionHandlers.getUnexpectedExceptionHandler(channel, true);
ErrorMessage error = ErrorMessage.fromException(t, handler);
error.setStreamId(request.getStreamId());
error.setWarnings(ClientWarn.instance.getWarnings());
return error;
return decorateResponse(ErrorMessage.fromException(t, request.getStreamId(), handler), request);
}
finally
{
Expand All @@ -485,6 +482,14 @@ void processRequest(Channel channel, Message.Request request, FlushItemConverter
flush(toFlush);
}

private static Message.Response decorateResponse(Message.Response response, Message.Request request)
{
assert response != null;
response.setStreamId(request.getStreamId());

@belliottsmith belliottsmith Jul 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If we're setting it here, I'm still unsure why we're setting it at all callers to ErrorMessage as well...

@frankgh frankgh Jul 15, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This method is not the only sink for responses, and that's why we still need to set it in the callers to ErrorMessage . For example, the callsites for ExceptionHandlers:81, PreV5Handlers:342, InitialConnectionHandler:134 / 163, and Dispatcher:115 do not flow through processRequest in the Dispatcher, and the stream ID would not be set. I'd prefer to have the redundancy here since it is a harmless operation.

@belliottsmith belliottsmith Jul 16, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The first two have no streamId. I suspect that InitialConnectionHandler must have a zero streamId as well. I am not sure it is even correct to set it there - it isn't set currently, have you checked the change is valid? That is, it's possible the streamId relates to work to be done after the connection is established, and the error message is meant to be at the connection level? I don't know, but since it is a change in behaviour we should check. Either way, we close the connection immediately in the first case, so we are not at risk of this bug, and we probably should close the connection in the second case as well since it is an unexpected message on startup so the connection is in a corrupted state. I would prefer to simply do this for these two cases anyway, since it's safer without extra validation (no possible change to behaviour, and we won't be at risk of the bug).

That leaves only the Dispatcher case, and you're right that we should address this in a uniform manner. However, this suggests to me that calls to FlushItemConverter should be forced to go through a method that performs this work for us, since this is the final step before writing the reply to the wire, so captures all normal replies. Ideally, we would do something similar for the out-of-band flushes in InitialConnectionHandler and ExceptionHandlers, so that there are at most two routes to send a message. This might be annoying to have the compiler enforce, but I think it would be OK to settle for a uniform invoker of toFlushItem, especially if we are diligent about only setting streamId when we know it, so that we will fail if we hit UNSET_STREAM_ID.

Regarding redundancy: I am not fond of redundancy for its own sake - our goal should always be to make sure the code is well structured so that we must use it correctly, and so we can have confidence that all edge cases are correctly handled (and future edge cases should not be easily added). Redundancy of the kind discussed here implies we aren't confident we can do that, which suggests we need to improve our abstractions instead.

This particular redundancy also has downsides, as it has lead to the use of NO_REQUEST_STREAM_ID in one place where it seems to weaken the protection we are introducing with UNSET_STREAM_ID, because we're enforcing that it be supplied in at least one place where it doesn't seem to make semantic sense.

That said, I won't push hard on this.

response.setWarnings(ClientWarn.instance.getWarnings());

@maedhroz maedhroz Jul 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this is too late to set warnings. The finally that resets warnings happens before this. The following test passes again if we move the warning setting back to the pre-patch location:

@Test
    public void warningsAreDeliveredToClient() throws Throwable
    {
        try (Cluster cluster = init(Cluster.build().withNodes(1)
                                           .withConfig(config -> config.with(GOSSIP, NETWORK, NATIVE_PROTOCOL)
                                                                       .set("tombstone_warn_threshold", 0)).start()))
        {
            cluster.schemaChange("CREATE TABLE " + KEYSPACE + ".tbl (pk int PRIMARY KEY, v int)");
            // insert a row so a tombstone is generated
            cluster.coordinator(1).execute("INSERT INTO " + KEYSPACE + ".tbl (pk, v) VALUES (1, 1)", ConsistencyLevel.ONE);
            cluster.coordinator(1).execute("DELETE v FROM " + KEYSPACE + ".tbl WHERE pk = 1", ConsistencyLevel.ONE);

            String address = cluster.get(1).config().broadcastAddress().getAddress().getHostAddress();
            int port = cluster.get(1).callOnInstance(() -> getNativeTransportPort());

            try (SimpleClient client = SimpleClient.builder(address, port)
                                                   .protocolVersion(ProtocolVersion.V5)
                                                   .build().connect(false))
            {
                QueryMessage query = new QueryMessage("SELECT * FROM " + KEYSPACE + ".tbl WHERE pk = 1", QueryOptions.DEFAULT);
                Message.Response response = client.execute(query);

                Assert.assertNotNull("Expected warnings on response, got none", response.getWarnings());
                Assert.assertFalse("Expected non-empty warnings list", response.getWarnings().isEmpty());
            }
        }
    }

@maedhroz maedhroz Jul 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actually, I think ClientWarningsTest already catches this.

@frankgh frankgh Jul 16, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is now fixed after the last commit

return response;
}

private void flush(FlushItem<?> item)
{
EventLoop loop = item.channel.eventLoop();
Expand Down
4 changes: 3 additions & 1 deletion src/java/org/apache/cassandra/transport/ExceptionHandlers.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,9 @@ public void exceptionCaught(final ChannelHandlerContext ctx, Throwable cause)
if (ctx.channel().isOpen())
{
Predicate<Throwable> handler = getUnexpectedExceptionHandler(ctx.channel(), false);
ErrorMessage errorMessage = ErrorMessage.fromException(cause, handler);
// No request in scope at the channel level; a WrappedException cause carries the frame's
// stream id and overrides this fallback, otherwise the channel is torn down for fatal errors.
ErrorMessage errorMessage = ErrorMessage.fromException(cause, Message.NO_REQUEST_STREAM_ID, handler);
Envelope response = errorMessage.encode(version);
FrameEncoder.Payload payload = allocator.allocate(true, CQLMessageHandler.envelopeSize(response.header));
try
Expand Down
5 changes: 3 additions & 2 deletions src/java/org/apache/cassandra/transport/InitialConnectionHandler.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -131,7 +131,7 @@ protected void decode(ChannelHandlerContext ctx, ByteBuf buffer, List<Object> li
if (null == cause)
cause = new ServerError("Unexpected error establishing connection");
logger.warn("Writing response to STARTUP failed, unable to configure pipeline", cause);
ErrorMessage error = ErrorMessage.fromException(cause);
ErrorMessage error = ErrorMessage.fromException(cause, inbound.header.streamId);
Envelope response = error.encode(inbound.header.version);
ChannelPromise closeChannel = AsyncChannelPromise.withListener(ctx, f -> ctx.close());
ctx.writeAndFlush(response, closeChannel);
Expand Down Expand Up @@ -161,7 +161,8 @@ protected void decode(ChannelHandlerContext ctx, ByteBuf buffer, List<Object> li
ErrorMessage error =
ErrorMessage.fromException(
new ProtocolException(String.format("Unexpected message %s, expecting STARTUP or OPTIONS",
inbound.header.type)));
inbound.header.type)),
inbound.header.streamId);
outbound = error.encode(inbound.header.version);
ctx.writeAndFlush(outbound);
}
Expand Down
13 changes: 13 additions & 0 deletions src/java/org/apache/cassandra/transport/Message.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,17 @@ public abstract class Message
{
protected static final Logger logger = LoggerFactory.getLogger(Message.class);

/** Sentinel default for a {@link Response}'s stream id; must be overwritten before {@link #encode} (asserted there). */
public static final int UNSET_STREAM_ID = Integer.MIN_VALUE;
/**
* Stream id used for channel-level errors that have no associated request (e.g. a raw protocol,
* SSL, or other unexpected failure caught in a pipeline {@code exceptionCaught} handler, where no
* request frame - and therefore no stream id - is in scope). A {@link ErrorMessage.WrappedException} cause, when
* present, carries the originating frame's stream id and overrides this value in
* {@link ErrorMessage#fromException(Throwable, int, com.google.common.base.Predicate)}.
*/
public static final int NO_REQUEST_STREAM_ID = 0;

public interface Codec<M extends Message> extends CBCodec<M> {}

public enum Direction
Expand Down Expand Up @@ -314,6 +325,7 @@ public static abstract class Response extends Message
protected Response(Type type)
{
super(type);
setStreamId(UNSET_STREAM_ID); // must be overwritten before encode

if (type.direction != Direction.RESPONSE)
throw new IllegalArgumentException();
Expand Down Expand Up @@ -430,6 +442,7 @@ public Envelope encode(ProtocolVersion version)
if (responseVersion.isBeta())
flags = Flag.add(flags, Flag.USE_BETA);

assert getStreamId() != UNSET_STREAM_ID : "Response streamId was never set: " + this;

@belliottsmith belliottsmith Jul 16, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This should not be an assertion, it should be a normal check and throw.

return Envelope.create(type, getStreamId(), responseVersion, flags, body);
}
catch (Throwable e)
Expand Down
4 changes: 3 additions & 1 deletion src/java/org/apache/cassandra/transport/PreV5Handlers.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -337,7 +337,9 @@ public void exceptionCaught(final ChannelHandlerContext ctx, Throwable cause)
if (ctx.channel().isOpen())
{
Predicate<Throwable> handler = ExceptionHandlers.getUnexpectedExceptionHandler(ctx.channel(), false);
ErrorMessage errorMessage = ErrorMessage.fromException(cause, handler);
// No request in scope at the channel level; a WrappedException cause carries the frame's
// stream id and overrides this fallback, otherwise the channel is torn down for fatal errors.
ErrorMessage errorMessage = ErrorMessage.fromException(cause, Message.NO_REQUEST_STREAM_ID, handler);

@belliottsmith belliottsmith Jul 16, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If we expect this to be overridden, why are we specifying NO_REQUEST_STREAM_ID, rather than UNSET_STREAM_ID?

@frankgh frankgh Jul 16, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There's one case where the negotiation fails during the protocol negotiation (app layer negotiation). This causes the org.apache.cassandra.transport.PreV5Handlers.ExceptionHandler#exceptionCaught code path above. And errorMessage.encode will hit the assertion in org.apache.cassandra.transport.Message#encode. If we set the stream Id to UNSET_STREAM_ID here, this would trigger the assertion error. So we'll need to do special handling for that case. This can be showcased in the org.apache.cassandra.transport.CQLConnectionTest#handleErrorDuringNegotiation test.

@frankgh frankgh Jul 16, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ok, I spent some time looking at this and I think we can always get a stream id except for one case. I will push a commit with these changes in a bit.

ChannelFuture future = ctx.writeAndFlush(errorMessage.encode(getConnectionVersion(ctx)));
// On protocol exception, close the channel as soon as the message have been sent.
// Most cases of PE are wrapped so the type check below is expected to fail more often than not.
Expand Down
2 changes: 1 addition & 1 deletion src/java/org/apache/cassandra/transport/messages/AuthResponse.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,6 @@ protected Response execute(QueryState queryState, Dispatcher.RequestTime request
{
return new AuthChallenge(challenge);
}
});
}, getStreamId());
}
}
7 changes: 5 additions & 2 deletions src/java/org/apache/cassandra/transport/messages/AuthUtil.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,10 +53,13 @@ public final class AuthUtil
* @param messageToSendBasedOnNegotiation Determines what response to return on based on whether sasl negotiation
* is complete (1st parameter) and the challenege token returned from the
* negotiator (2nd parameter).
* @param streamId The stream id of the request being handled, used to stamp any error
* response so it is routed back to the correct in-flight request.
* @return the response to send back to the client.
*/
static Response handleLogin(Connection connection, QueryState queryState, byte[] token,
BiFunction<Boolean, byte[], Response> messageToSendBasedOnNegotiation)
BiFunction<Boolean, byte[], Response> messageToSendBasedOnNegotiation,
int streamId)
{
IAuthenticator.SaslNegotiator negotiator = ((ServerConnection) connection).getSaslNegotiator(queryState);
try
Expand Down Expand Up @@ -87,7 +90,7 @@ static Response handleLogin(Connection connection, QueryState queryState, byte[]
{
ClientMetrics.instance.markAuthFailure(negotiator.getAuthenticationMode());
AuthEvents.instance.notifyAuthFailure(queryState, e);
return ErrorMessage.fromException(e);
return ErrorMessage.fromException(e, streamId);
}
}

Expand Down
2 changes: 1 addition & 1 deletion src/java/org/apache/cassandra/transport/messages/BatchMessage.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -236,7 +236,7 @@ protected Message.Response execute(QueryState state, Dispatcher.RequestTime requ
{
QueryEvents.instance.notifyBatchFailure(prepared, batchType, queryOrIdList, values, options, state, e);
JVMStabilityInspector.inspectThrowable(e);
return ErrorMessage.fromException(e);
return ErrorMessage.fromException(e, getStreamId());
}
}

Expand Down
11 changes: 6 additions & 5 deletions src/java/org/apache/cassandra/transport/messages/ErrorMessage.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -429,20 +429,21 @@ private ErrorMessage(TransportException error, int streamId)
setStreamId(streamId);
}

public static ErrorMessage fromException(Throwable e)
public static ErrorMessage fromException(Throwable e, int streamId)
{
return fromException(e, null);
return fromException(e, streamId, null);
}

/**
* @param e the exception
* @param streamId the stream id of the request this error responds to, so the error frame is routed back
* to the correct in-flight request. A {@link WrappedException} cause overrides this with the
* stream id recovered from the frame header.
* @param unexpectedExceptionHandler a callback for handling unexpected exceptions. If null, or if this
* returns false, the error is logged at ERROR level via sl4fj
*/
public static ErrorMessage fromException(Throwable e, Predicate<Throwable> unexpectedExceptionHandler)
public static ErrorMessage fromException(Throwable e, int streamId, Predicate<Throwable> unexpectedExceptionHandler)
{
int streamId = 0;

// Netty will wrap exceptions during decoding in a CodecException. If the cause was one of our ProtocolExceptions
// or some other internal exception, extract that and use it.
if (e instanceof CodecException)
Expand Down
2 changes: 1 addition & 1 deletion src/java/org/apache/cassandra/transport/messages/ExecuteMessage.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -207,7 +207,7 @@ else if (options.skipMetadata())
{
QueryEvents.instance.notifyExecuteFailure(prepared, options, state, e);
JVMStabilityInspector.inspectThrowable(e);
return ErrorMessage.fromException(e);
return ErrorMessage.fromException(e, getStreamId());
}
}

Expand Down
2 changes: 1 addition & 1 deletion src/java/org/apache/cassandra/transport/messages/PrepareMessage.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -134,7 +134,7 @@ protected Message.Response execute(QueryState state, Dispatcher.RequestTime requ
{
QueryEvents.instance.notifyPrepareFailure(null, query, state, e);
JVMStabilityInspector.inspectThrowable(e);
return ErrorMessage.fromException(e);
return ErrorMessage.fromException(e, getStreamId());
}
}

Expand Down
2 changes: 1 addition & 1 deletion src/java/org/apache/cassandra/transport/messages/QueryMessage.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -129,7 +129,7 @@ protected Message.Response execute(QueryState state, Dispatcher.RequestTime requ
JVMStabilityInspector.inspectThrowable(e);
if (!((e instanceof RequestValidationException) || (e instanceof RequestExecutionException)))
logger.error("Unexpected error during query", e);
return ErrorMessage.fromException(e);
return ErrorMessage.fromException(e, getStreamId());
}
}

Expand Down
2 changes: 1 addition & 1 deletion src/java/org/apache/cassandra/transport/messages/StartupMessage.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -158,7 +158,7 @@ else if (compression.equals("lz4"))
// authentication, in this case we can just go through the traditional auth flow.
return authenticator.getAuthenticateMessage(clientState);
}
});
}, getStreamId());
}
}
return authenticator.getAuthenticateMessage(clientState);
Expand Down
Loading