Merge pull request #104 from scalecube/fix-sending-complete-signal

Fix sending complete signal

Author: segabriel

services-gateway-netty/src/main/java/io/scalecube/services/gateway/GatewaySessionHandler.java

@@ -23,7 +23,7 @@ public interface GatewaySessionHandler<M> {
    * @param req request message (not null)
    * @return message
    */
-  default M mapMessage(GatewaySession session, M req) {
+  default M mapMessage(GatewaySession session, M req, Context context) {
     return req;
   }
 

services-gateway-netty/src/main/java/io/scalecube/services/gateway/rsocket/RSocketGatewayAcceptor.java

@@ -12,6 +12,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import reactor.core.publisher.Mono;
+import reactor.util.context.Context;
 
 public class RSocketGatewayAcceptor implements SocketAcceptor {
 
@@ -46,7 +47,10 @@ public Mono<RSocket> accept(ConnectionSetupPayload setup, RSocket rsocket) {
     ServiceMessageCodec messageCodec = new ServiceMessageCodec(headersCodec);
     final RSocketGatewaySession gatewaySession =
         new RSocketGatewaySession(
-            serviceCall, metrics, messageCodec, gatewaySessionHandler::mapMessage);
+            serviceCall,
+            metrics,
+            messageCodec,
+            (session, req) -> gatewaySessionHandler.mapMessage(session, req, Context.empty()));
     gatewaySessionHandler.onSessionOpen(gatewaySession);
     rsocket
         .onClose()

services-gateway-netty/src/main/java/io/scalecube/services/gateway/ws/WebsocketGatewayAcceptor.java

@@ -78,7 +78,7 @@ private Mono<GatewayMessage> onRequest(
         .flatMap(msg -> onCancel(session, msg))
         .map(msg -> validateSid(session, (GatewayMessage) msg))
         .map(this::validateQualifier)
-        .map(msg -> gatewayHandler.mapMessage(session, msg))
+        .map(msg -> gatewayHandler.mapMessage(session, msg, context))
         .doOnNext(request -> onMessage(session, request, context))
         .doOnError(
             th -> {
@@ -108,10 +108,10 @@ private void onMessage(WebsocketGatewaySession session, GatewayMessage request,
             .orElse(serviceStream)
             .map(response -> prepareResponse(sid, response, receivedError))
             .doOnNext(response -> metrics.markServiceResponse())
-            .doFinally(signalType -> session.dispose(sid))
             .flatMap(session::send)
             .doOnError(th -> onError(session, request, th, context))
             .doOnComplete(() -> onComplete(session, request, receivedError, context))
+            .doFinally(signalType -> session.dispose(sid))
             .subscriberContext(context)
             .subscribe();
 

services-gateway-tests/src/main/java/io/scalecube/services/testservice/SecuredRsGwGatewaySessionHandler.java

@@ -5,6 +5,7 @@
 import io.scalecube.services.gateway.GatewaySessionHandler;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import reactor.util.context.Context;
 
 public class SecuredRsGwGatewaySessionHandler implements GatewaySessionHandler<ServiceMessage> {
   private static final Logger LOGGER =
@@ -16,7 +17,7 @@ public SecuredRsGwGatewaySessionHandler(AuthRegistry authRegistry) {
   }
 
   @Override
-  public ServiceMessage mapMessage(GatewaySession session, ServiceMessage req) {
+  public ServiceMessage mapMessage(GatewaySession session, ServiceMessage req, Context context) {
     return ServiceMessage.from(req).header(AuthRegistry.SESSION_ID, session).build();
   }
 

services-gateway-tests/src/main/java/io/scalecube/services/testservice/SecuredWsGwGatewaySessionManager.java

@@ -5,6 +5,7 @@
 import io.scalecube.services.gateway.ws.GatewayMessage;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import reactor.util.context.Context;
 
 public class SecuredWsGwGatewaySessionManager implements GatewaySessionHandler<GatewayMessage> {
   private static final Logger LOGGER =
@@ -16,7 +17,7 @@ public SecuredWsGwGatewaySessionManager(AuthRegistry authRegistry) {
   }
 
   @Override
-  public GatewayMessage mapMessage(GatewaySession session, GatewayMessage req) {
+  public GatewayMessage mapMessage(GatewaySession session, GatewayMessage req, Context context) {
     return GatewayMessage.from(req).header(AuthRegistry.SESSION_ID, session).build();
   }
 

services-gateway-tests/src/test/java/io/scalecube/services/gateway/TestGatewaySessionHandler.java

@@ -1,6 +1,7 @@
 package io.scalecube.services.gateway;
 
 import java.util.concurrent.CountDownLatch;
+import reactor.util.context.Context;
 
 public class TestGatewaySessionHandler implements GatewaySessionHandler {
 
@@ -9,7 +10,7 @@ public class TestGatewaySessionHandler implements GatewaySessionHandler {
   public final CountDownLatch disconnLatch = new CountDownLatch(1);
 
   @Override
-  public Object mapMessage(GatewaySession s, Object req) {
+  public Object mapMessage(GatewaySession s, Object req, Context context) {
     msgLatch.countDown();
     return req;
   }