diff --git a/spring-messaging/src/main/java/org/springframework/messaging/rsocket/annotation/support/MessagingRSocket.java b/spring-messaging/src/main/java/org/springframework/messaging/rsocket/annotation/support/MessagingRSocket.java index f2f1da621dd..19a6eb2f68b 100644 --- a/spring-messaging/src/main/java/org/springframework/messaging/rsocket/annotation/support/MessagingRSocket.java +++ b/spring-messaging/src/main/java/org/springframework/messaging/rsocket/annotation/support/MessagingRSocket.java @@ -103,8 +103,9 @@ class MessagingRSocket implements RSocket { * @return completion handle for success or error */ public Mono handleConnectionSetupPayload(ConnectionSetupPayload payload) { - // frameDecoder does not apply to connectionSetupPayload - // so retain here since handle expects it. + // RSocket Java releases its own ConnectionSetupPayload reference + // For other frames, PayloadDecoder increases the refCount, and handles must releases it; + // retain here to match the release in handle (in retainDataAndReleasePayload) payload.retain(); return handle(payload, FrameType.SETUP); } @@ -187,32 +188,40 @@ class MessagingRSocket implements RSocket { private MessageHeaders createHeaders( Payload payload, FrameType frameType, @Nullable AtomicReference> responseRef) { - MessageHeaderAccessor headers = new MessageHeaderAccessor(); - headers.setLeaveMutable(true); + try { + MessageHeaderAccessor headers = new MessageHeaderAccessor(); + headers.setLeaveMutable(true); - Map metadataValues = this.metadataExtractor.extract(payload, this.metadataMimeType); + Map metadataValues = this.metadataExtractor.extract(payload, this.metadataMimeType); - metadataValues.putIfAbsent(MetadataExtractor.ROUTE_KEY, ""); - for (Map.Entry entry : metadataValues.entrySet()) { - if (entry.getKey().equals(MetadataExtractor.ROUTE_KEY)) { - RouteMatcher.Route route = this.routeMatcher.parseRoute((String) entry.getValue()); - headers.setHeader(DestinationPatternsMessageCondition.LOOKUP_DESTINATION_HEADER, route); + metadataValues.putIfAbsent(MetadataExtractor.ROUTE_KEY, ""); + for (Map.Entry entry : metadataValues.entrySet()) { + if (entry.getKey().equals(MetadataExtractor.ROUTE_KEY)) { + RouteMatcher.Route route = this.routeMatcher.parseRoute((String) entry.getValue()); + headers.setHeader(DestinationPatternsMessageCondition.LOOKUP_DESTINATION_HEADER, route); + } + else { + headers.setHeader(entry.getKey(), entry.getValue()); + } } - else { - headers.setHeader(entry.getKey(), entry.getValue()); + + headers.setContentType(this.dataMimeType); + headers.setHeader(RSocketFrameTypeMessageCondition.FRAME_TYPE_HEADER, frameType); + headers.setHeader(RSocketRequesterMethodArgumentResolver.RSOCKET_REQUESTER_HEADER, this.requester); + if (responseRef != null) { + headers.setHeader(RSocketPayloadReturnValueHandler.RESPONSE_HEADER, responseRef); } - } + headers.setHeader(HandlerMethodReturnValueHandler.DATA_BUFFER_FACTORY_HEADER, + this.strategies.dataBufferFactory()); - headers.setContentType(this.dataMimeType); - headers.setHeader(RSocketFrameTypeMessageCondition.FRAME_TYPE_HEADER, frameType); - headers.setHeader(RSocketRequesterMethodArgumentResolver.RSOCKET_REQUESTER_HEADER, this.requester); - if (responseRef != null) { - headers.setHeader(RSocketPayloadReturnValueHandler.RESPONSE_HEADER, responseRef); + return headers.getMessageHeaders(); + } + catch (Throwable ex) { + if (payload.refCnt() > 0) { + payload.release(); + } + throw ex; } - headers.setHeader(HandlerMethodReturnValueHandler.DATA_BUFFER_FACTORY_HEADER, - this.strategies.dataBufferFactory()); - - return headers.getMessageHeaders(); } }