Emit PgNotification on a duplicated context

Currently, notifications are emitted on the context captured when the connection is created.
If a PgSubscriber was created while running on a DuplicatedContext, this would be context of emissions.
This is wrong and could cause local data (e.g. tracing) to be still present when notification handlers are executed.

This change consists in unwrapping the connection context at creation time, and then duplicating it before emitting notifications.

Signed-off-by: Thomas Segismont <tsegismont@gmail.com>

Author: tsegismont

vertx-pg-client/src/main/java/io/vertx/pgclient/impl/PgConnectionImpl.java

@@ -28,9 +28,9 @@
 import io.vertx.pgclient.impl.codec.NoticeResponse;
 import io.vertx.pgclient.impl.codec.TxFailedEvent;
 import io.vertx.pgclient.spi.PgDriver;
-import io.vertx.sqlclient.spi.connection.Connection;
 import io.vertx.sqlclient.codec.SocketConnectionBase;
 import io.vertx.sqlclient.internal.SqlConnectionBase;
+import io.vertx.sqlclient.spi.connection.Connection;
 
 public class PgConnectionImpl extends SqlConnectionBase<PgConnectionImpl> implements PgConnection  {
 
@@ -67,10 +67,11 @@ public void handleEvent(Object event) {
       Handler<PgNotification> handler = notificationHandler;
       if (handler != null) {
         Notification notification = (Notification) event;
-        handler.handle(new PgNotification()
+        PgNotification pgNotification = new PgNotification()
           .setChannel(notification.getChannel())
           .setProcessId(notification.getProcessId())
-          .setPayload(notification.getPayload()));
+          .setPayload(notification.getPayload());
+        context.duplicate().emit(pgNotification, handler);
       }
     } else if (event instanceof NoticeResponse) {
       Handler<PgNotice> handler = noticeHandler;
@@ -94,8 +95,7 @@ public void handleEvent(Object event) {
         .setDataType(noticeEvent.getDataType())
         .setConstraint(noticeEvent.getConstraint());
       if (handler != null) {
-        handler.handle(notice
-        );
+        context.duplicate().emit(notice, handler);
       } else {
         notice.log(SocketConnectionBase.logger);
       }

vertx-pg-client/src/test/java/io/vertx/tests/pgclient/PubSubTest.java

@@ -18,6 +18,7 @@
 
 import io.vertx.core.Vertx;
 import io.vertx.core.buffer.Buffer;
+import io.vertx.core.internal.ContextInternal;
 import io.vertx.ext.unit.Async;
 import io.vertx.ext.unit.TestContext;
 import io.vertx.pgclient.PgConnectOptions;
@@ -32,9 +33,12 @@
 
 import java.util.Arrays;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 
+import static io.vertx.core.internal.ContextInternal.LOCAL_MAP;
+
 public class PubSubTest extends PgTestBase {
 
   Vertx vertx;
@@ -361,4 +365,52 @@ public void testNoticedRaised(TestContext ctx) {
       }));
     }));
   }
+
+  @Test
+  public void testSubscriberEmitsOnDuplicatedContext(TestContext ctx) {
+    String channelName = "test_channel";
+    String quotedChannelName = "\"" + channelName + "\"";
+
+    // Create a duplicated context with local data BEFORE creating the subscriber
+    ContextInternal originalContext = (ContextInternal) vertx.getOrCreateContext();
+    ContextInternal duplicatedContext = originalContext.duplicate();
+    duplicatedContext.getLocal(LOCAL_MAP, ConcurrentHashMap::new).put("foo", "bar");
+
+    Async connectLatch = ctx.async();
+
+    duplicatedContext.runOnContext(v -> {
+      ctx.assertEquals("bar", ContextInternal.current().getLocal(LOCAL_MAP, ConcurrentHashMap::new).get("foo"));
+
+      subscriber = PgSubscriber.subscriber(vertx, options);
+      subscriber.connect().onComplete(v2 -> connectLatch.complete());
+    });
+
+    connectLatch.awaitSuccess(10000);
+
+    Async notificationLatch = ctx.async();
+
+    duplicatedContext.runOnContext(v -> {
+      ctx.assertEquals("bar", ContextInternal.current().getLocal(LOCAL_MAP, ConcurrentHashMap::new).get("foo"));
+
+      PgChannel channel = subscriber.channel(channelName);
+      channel.handler(notif -> {
+        ContextInternal currentContext = ContextInternal.current();
+        ctx.assertTrue(currentContext.isDuplicate());
+        ctx.assertNotEquals(duplicatedContext, currentContext);
+        ctx.assertNull(currentContext.getLocal(LOCAL_MAP, ConcurrentHashMap::new).get("foo"));
+        ctx.assertEquals("msg", notif);
+        notificationLatch.countDown();
+      });
+    });
+
+    vertx.setTimer(100, t -> {
+      subscriber
+        .actualConnection()
+        .query("NOTIFY " + quotedChannelName + ", 'msg'")
+        .execute()
+        .onComplete(ctx.asyncAssertSuccess());
+    });
+
+    notificationLatch.awaitSuccess(10000);
+  }
 }

vertx-sql-client/src/main/java/io/vertx/sqlclient/internal/SqlConnectionBase.java

@@ -18,26 +18,26 @@
 package io.vertx.sqlclient.internal;
 
 import io.vertx.core.*;
-import io.vertx.core.spi.metrics.ClientMetrics;
-import io.vertx.core.spi.tracing.VertxTracer;
 import io.vertx.core.internal.ContextInternal;
 import io.vertx.core.internal.PromiseInternal;
+import io.vertx.core.spi.metrics.ClientMetrics;
+import io.vertx.core.spi.tracing.VertxTracer;
 import io.vertx.sqlclient.PrepareOptions;
 import io.vertx.sqlclient.PreparedStatement;
 import io.vertx.sqlclient.SqlConnection;
 import io.vertx.sqlclient.Transaction;
 import io.vertx.sqlclient.impl.PreparedStatementBase;
 import io.vertx.sqlclient.impl.TransactionImpl;
-import io.vertx.sqlclient.spi.connection.ConnectionContext;
-import io.vertx.sqlclient.spi.protocol.CommandBase;
-import io.vertx.sqlclient.spi.connection.Connection;
-import io.vertx.sqlclient.spi.protocol.PrepareStatementCommand;
-import io.vertx.sqlclient.spi.protocol.QueryCommandBase;
 import io.vertx.sqlclient.impl.pool.SqlConnectionPool;
 import io.vertx.sqlclient.impl.tracing.QueryReporter;
-import io.vertx.sqlclient.spi.connection.ConnectionFactory;
 import io.vertx.sqlclient.spi.DatabaseMetadata;
 import io.vertx.sqlclient.spi.Driver;
+import io.vertx.sqlclient.spi.connection.Connection;
+import io.vertx.sqlclient.spi.connection.ConnectionContext;
+import io.vertx.sqlclient.spi.connection.ConnectionFactory;
+import io.vertx.sqlclient.spi.protocol.CommandBase;
+import io.vertx.sqlclient.spi.protocol.PrepareStatementCommand;
+import io.vertx.sqlclient.spi.protocol.QueryCommandBase;
 
 /**
  * @author <a href="mailto:julien@julienviet.com">Julien Viet</a>
@@ -54,7 +54,7 @@ public class SqlConnectionBase<C extends SqlConnectionBase<C>> extends SqlClient
 
   public SqlConnectionBase(ContextInternal context, ConnectionFactory factory, Connection conn, Driver driver) {
     super(driver);
-    this.context = context;
+    this.context = context.unwrap();
     this.factory = factory;
     this.conn = conn;
   }