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 3 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
3 commits
Select commit Hold shift + click to select a range
4f1f2d9
Ensure a Message's Response streamId is always set
frankgh Jul 13, 2026
79af4a6
Cleanup after manual review
frankgh Jul 24, 2026
046f959
Fix
frankgh Jul 25, 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
20 changes: 11 additions & 9 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 @@ -106,7 +106,7 @@ public class CQLMessageHandler<M extends Message> extends AbstractMessageHandler

interface MessageConsumer<M extends Message>
{
void dispatch(Channel channel, M message, Dispatcher.FlushItemConverter toFlushItem, Overload backpressure);
<P> void dispatch(Channel channel, M message, Dispatcher.FlushItemConverter<P> toFlushItem, P param, Overload backpressure);
boolean hasQueueCapacity();
}

Expand Down Expand Up @@ -389,7 +389,7 @@ protected boolean processRequest(Envelope request, Overload backpressure)
try
{
message = messageDecoder.decode(channel, request);
dispatcher.dispatch(channel, message, this::toFlushItem, backpressure);
dispatcher.dispatch(channel, message, CQLMessageHandler::toFlushItem, this, backpressure);

// sucessfully delivered a CQL message to the execution
// stage, so reset the counter of consecutive errors
Expand Down Expand Up @@ -485,7 +485,8 @@ private Framed toFlushItem(Channel channel, Message.Request request, Message.Res
// The Dispatcher will call this to obtain the FlushItem to enqueue with its Flusher once
// a dispatched request has been processed.

Envelope responseFrame = response.encode(request.getSource().header.version);
Envelope.Header header = request.getSource().header;
Envelope responseFrame = response.encode(header.version, header.streamId);
int responseSize = envelopeSize(responseFrame.header);
ClientMessageSizeMetrics.bytesSent.inc(responseSize);
ClientMessageSizeMetrics.bytesSentPerResponse.update(responseSize);
Expand All @@ -500,9 +501,9 @@ private Framed toFlushItem(Channel channel, Message.Request request, Message.Res
@Override
public void cleanup(Flusher.FlushItem<Envelope> flushItem)
{
release(flushItem.request.header);
flushItem.request.release();
flushItem.response.release();
release(flushItem.requestEnvelope.header);
flushItem.requestEnvelope.release();
flushItem.responseEnvelope.release();
}

private void release(Envelope.Header header)
Expand All @@ -524,8 +525,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 +544,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
49 changes: 24 additions & 25 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 @@ -93,10 +93,9 @@ public class Dispatcher implements CQLMessageHandler.MessageConsumer<Message.Req
* The instances of these FlushItem subclasses are specialized to release resources in the
* right way for the specific pipeline that produced them.
*/
// TODO parameterize with FlushItem subclass
interface FlushItemConverter
public interface FlushItemConverter<P>
{
FlushItem<?> toFlushItem(Channel channel, Message.Request request, Message.Response response);
FlushItem<?> toFlushItem(P param, Channel channel, Message.Request request, Message.Response response);
}

public Dispatcher(boolean useLegacyFlusher)
Expand All @@ -105,18 +104,17 @@ public Dispatcher(boolean useLegacyFlusher)
}

@Override
public void dispatch(Channel channel, Message.Request request, FlushItemConverter forFlusher, Overload backpressure)
public <P> void dispatch(Channel channel, Message.Request request, FlushItemConverter<P> forFlusher, P param, Overload backpressure)
{
if (!request.connection().getTracker().isRunning())
{
// 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.fromTransportException(new OverloadedException("Server is shutting down"));
response.setWarnings(ClientWarn.instance.getWarnings());
response.attach(request.connection);
FlushItem<?> toFlush = forFlusher.toFlushItem(channel, request, response);
FlushItem<?> toFlush = forFlusher.toFlushItem(param, channel, request, response);
flush(toFlush);
return;
}
Expand All @@ -128,7 +126,7 @@ public void dispatch(Channel channel, Message.Request request, FlushItemConverte
// Importantly, the authExecutor will handle the AUTHENTICATE message which may be CPU intensive.
LocalAwareExecutorPlus executor = isAuthQuery ? authExecutor : requestExecutor;

executor.submit(new RequestProcessor(channel, request, forFlusher, backpressure));
executor.submit(new RequestProcessor<>(channel, request, forFlusher, param, backpressure));
ClientMetrics.instance.markRequestDispatched();
}

Expand Down Expand Up @@ -293,28 +291,30 @@ public long timeSpentInQueueNanos()
* is the only way we can keep it not wrapped into a callable on SEPExecutor submission path. And we need this
* functionality for tracking time purposes.
*/
public class RequestProcessor implements DebuggableTask.RunnableDebuggableTask
public class RequestProcessor<P> implements DebuggableTask.RunnableDebuggableTask
{
private final Channel channel;
private final Message.Request request;
private final FlushItemConverter forFlusher;
private final FlushItemConverter<P> forFlusher;
private final P flusherParam;
private final Overload backpressure;

private volatile long startTimeNanos;

public RequestProcessor(Channel channel, Message.Request request, FlushItemConverter forFlusher, Overload backpressure)
public RequestProcessor(Channel channel, Message.Request request, FlushItemConverter<P> forFlusher, P flusherParam, Overload backpressure)
{
this.channel = channel;
this.request = request;
this.forFlusher = forFlusher;
this.flusherParam = flusherParam;
this.backpressure = backpressure;
}

@Override
public void run()
{
startTimeNanos = MonotonicClock.Global.preciseTime.now();
processRequest(channel, request, forFlusher, backpressure, new RequestTime(request.createdAtNanos, startTimeNanos));
processRequest(channel, request, forFlusher, flusherParam, backpressure, new RequestTime(request.createdAtNanos, startTimeNanos));
}

@Override
Expand Down Expand Up @@ -374,7 +374,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.fromTransportException(new OverloadedException("Query timed out before it could start"));
}

if (connection.getVersion().isGreaterOrEqualTo(ProtocolVersion.V4))
Expand Down Expand Up @@ -434,8 +434,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 @@ -446,9 +444,10 @@ private static Message.Response processRequest(ServerConnection connection, Mess
*/
static Message.Response processRequest(Channel channel, Message.Request request, Overload backpressure, RequestTime requestTime)
{
Message.Response response = null;
try
{
return processRequest((ServerConnection) request.connection(), request, backpressure, requestTime);
response = processRequest((ServerConnection) request.connection(), request, backpressure, requestTime);
}
catch (Throwable t)
{
Expand All @@ -461,26 +460,26 @@ 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;
response = ErrorMessage.fromExceptionNoStreamId(t, handler);
}
finally
{
if (response != null)
response.setWarnings(ClientWarn.instance.getWarnings());
CoordinatorWarnings.reset();
CoordinatorWriteWarnings.reset();
ClientWarn.instance.resetWarnings();
}
return response;
}

/**
* Note: this method is not expected to execute on the netty event loop.
*/
void processRequest(Channel channel, Message.Request request, FlushItemConverter forFlusher, Overload backpressure, RequestTime requestTime)
<P> void processRequest(Channel channel, Message.Request request, FlushItemConverter<P> forFlusher, P param, Overload backpressure, RequestTime requestTime)
{
Message.Response response = processRequest(channel, request, backpressure, requestTime);
FlushItem<?> toFlush = forFlusher.toFlushItem(channel, request, response);
FlushItem<?> toFlush = forFlusher.toFlushItem(param, channel, request, response);
Message.logger.trace("Responding: {}, v={}", response, request.connection().getVersion());
flush(toFlush);
}
Expand Down Expand Up @@ -517,7 +516,7 @@ public static void shutdown()
* for delivering events to registered clients is dependent on protocol version and the configuration
* of the pipeline. For v5 and newer connections, the event message is encoded into an Envelope,
* wrapped in a FlushItem and then delivered via the pipeline's flusher, in a similar way to
* a Response returned from {@link #processRequest(Channel, Message.Request, FlushItemConverter, Overload, RequestTime)}.
* a Response returned from {@link #processRequest(Channel, Message.Request, FlushItemConverter, Object, Overload, RequestTime)}.
* It's worth noting that events are not generally fired as a direct response to a client request,
* so this flush item has a null request attribute. The dispatcher itself is created when the
* pipeline is first configured during protocol negotiation and is attached to the channel for
Expand All @@ -531,9 +530,9 @@ Consumer<EventMessage> eventDispatcher(final Channel channel,
final FrameEncoder.PayloadAllocator allocator)
{
return eventMessage -> flush(new FlushItem.Framed(channel,
eventMessage.encode(version),
eventMessage.encode(version, EventMessage.EVENT_MESSAGE_STREAM_ID), // -1 was set in EventMessage previously
null,
allocator,
f -> f.response.release()));
f -> f.responseEnvelope.release()));
}
}
55 changes: 45 additions & 10 deletions src/java/org/apache/cassandra/transport/Envelope.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -146,6 +146,11 @@ public static class Header
public final Message.Type type;
public final long bodySizeInBytes;

public static Header dummy(int streamId, Message.Type type)
{
return new Header(ProtocolVersion.CURRENT, 0, streamId, type, 0);
}

private Header(ProtocolVersion version, int flags, int streamId, Message.Type type, long bodySizeInBytes)
{
this.version = version;
Expand Down Expand Up @@ -258,7 +263,7 @@ HeaderExtractionResult extractHeader(ByteBuffer buffer)
// This throws a protocol exception if the version number is unsupported,
// the opcode is unknown or invalid flags are set for the version
version = ProtocolVersion.decode(versionNum, DatabaseDescriptor.getNativeTransportAllowOlderProtocols());
validateFlags(version, flags);
validateFlags(version, flags, streamId);
type = Message.Type.fromOpcode(opcode, direction);
return new HeaderExtractionResult.Success(new Header(version, flags, streamId, type, bodyLength));
}
Expand All @@ -272,6 +277,10 @@ HeaderExtractionResult extractHeader(ByteBuffer buffer)
// cause the channel to be closed.
return new HeaderExtractionResult.Error(e, streamId, bodyLength);
}
catch (ErrorMessage.WrappedException e)
{
return new HeaderExtractionResult.Error((ProtocolException) e.getCause(), e.getStreamId(), bodyLength);
}
}

public static abstract class HeaderExtractionResult
Expand Down Expand Up @@ -370,27 +379,52 @@ Envelope decode(ByteBuf buffer)
Message.Direction direction = Message.Direction.extractFromVersion(firstByte);
int versionNum = firstByte & PROTOCOL_VERSION_MASK;

ProtocolVersion version;
ProtocolVersion version = null;
ProtocolException protocolException = null;

try
{
version = ProtocolVersion.decode(versionNum, DatabaseDescriptor.getNativeTransportAllowOlderProtocols());
}
catch (ProtocolException e)
{
// Skip the remaining useless bytes. Otherwise the channel closing logic may try to decode again.
buffer.skipBytes(readableBytes);
throw e;
// defer throw to attempt to extract the stream id
protocolException = e;
}

// Wait until we have the complete header
if (readableBytes < Header.LENGTH)
{
if (protocolException != null)
{
// Skip the remaining useless bytes. Otherwise the channel closing logic may try to decode again.
buffer.skipBytes(readableBytes);
// The header is incomplete, so there is no stream id to recover. Wrap with the unset
// sentinel; the channel-level exception handler sees it has no routable stream id and
// closes the connection rather than emit an unroutable error frame.
throw protocolException;
}
return null;
}

int flags = buffer.getByte(idx++);
validateFlags(version, flags);

int streamId = buffer.getShort(idx);

if (protocolException != null)
{
// Protocol versions 1 and 2 use a shorter header with a single-byte stream id. Reading a
// 16-bit stream id from such a header splices the stream id byte together with the opcode
// byte and recovers a bogus id, routing the error to a stream the client never used
// (CASSANDRA-21508). A v1/v2 client that downgrades would then never see the error and time
// out, so recover the stream id using the attempted version's header layout.
int recoveredStreamId = versionNum < ProtocolVersion.V3.asInt() ? buffer.getByte(idx) : streamId;
// Skip the remaining useless bytes. Otherwise the channel closing logic may try to decode again.
buffer.skipBytes(readableBytes);
throw ErrorMessage.wrap(protocolException, recoveredStreamId);
}

validateFlags(version, flags, streamId);

idx += 2;

// This throws a protocol exceptions if the opcode is unknown
Expand Down Expand Up @@ -436,11 +470,12 @@ Envelope decode(ByteBuf buffer)
return new Envelope(new Header(version, flags, streamId, type, bodyLength), body);
}

private void validateFlags(ProtocolVersion version, int flags)
private void validateFlags(ProtocolVersion version, int flags, int streamId)
{
if (version.isBeta() && !Header.Flag.contains(flags, Header.Flag.USE_BETA))
throw new ProtocolException(String.format("Beta version of the protocol used (%s), but USE_BETA flag is unset", version),
version);
throw ErrorMessage.wrap(new ProtocolException(String.format("Beta version of the protocol used (%s), but USE_BETA flag is unset", version),
version),
streamId);
}

@Override
Expand Down
Loading