fix

hibernate · Feb 3, 2025 · 51372f1 · 51372f1
1 parent 4b1857b
commit 51372f1
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 5 deletions.
diff --git a/hibernate-reactive-core/src/main/java/org/hibernate/reactive/context/Context.java b/hibernate-reactive-core/src/main/java/org/hibernate/reactive/context/Context.java
@@ -9,6 +9,8 @@
 
 import java.util.concurrent.Executor;
 
+import io.vertx.core.spi.context.storage.ContextLocal;
+
 /**
  * Abstracts away from the Vert.x {@link io.vertx.core.Context}
  * object, enabling alternative strategies for associating state
@@ -53,5 +55,5 @@ public interface Context extends Executor, Service {
      *
      * @param <T> the type of thing we're storing in the context
      */
-    interface Key<T> {}
+    interface Key<T> extends ContextLocal<T> {}
 }
diff --git a/hibernate-reactive-core/src/main/java/org/hibernate/reactive/context/impl/VertxContext.java b/hibernate-reactive-core/src/main/java/org/hibernate/reactive/context/impl/VertxContext.java
@@ -36,9 +36,10 @@ public void injectServices(ServiceRegistryImplementor serviceRegistry) {
 
 	@Override
 	public <T> void put(Key<T> key, T instance) {
-		final io.vertx.core.Context context = Vertx.currentContext();
+		final ContextInternal context = currentContext();
 		if ( context != null ) {
 			if ( trace ) LOG.tracef( "Putting key,value in context: [%1$s, %2$s]", key, instance );
+			context.putLocal(  );
 			context.putLocal( key, instance );
 		}
 		else {
@@ -47,9 +48,13 @@ public <T> void put(Key<T> key, T instance) {
 		}
 	}
 
+	private static ContextInternal currentContext() {
+		return (ContextInternal) Vertx.currentContext();
+	}
+
 	@Override
 	public <T> T get(Key<T> key) {
-		final io.vertx.core.Context context = Vertx.currentContext();
+		final ContextInternal context = currentContext();
 		if ( context != null ) {
 			T local = context.getLocal( key );
 			if ( trace ) LOG.tracef( "Getting value %2$s from context for key %1$s", key, local  );
@@ -63,7 +68,7 @@ public <T> T get(Key<T> key) {
 
 	@Override
 	public void remove(Key<?> key) {
-		final io.vertx.core.Context context = Vertx.currentContext();
+		final ContextInternal context = currentContext();
 		if ( context != null ) {
 			boolean removed = context.removeLocal( key );
 			if ( trace ) LOG.tracef( "Key %s removed from context: %s", key, removed );
@@ -75,7 +80,7 @@ public void remove(Key<?> key) {
 
 	@Override
 	public void execute(Runnable runnable) {
-		final io.vertx.core.Context currentContext = Vertx.currentContext();
+		final io.vertx.core.Context currentContext = currentContext();
 		if ( currentContext == null ) {
 			if ( trace ) LOG.tracef( "Not in a Vert.x context, checking the VertxInstance service" );
 			final io.vertx.core.Context newContext = vertxInstance.getVertx().getOrCreateContext();