Create custom Vert.x contextual data storage

hibernate-reactive-core/src/main/java/org/hibernate/reactive/context/impl/ContextualDataStorage.java

@@ -0,0 +1,25 @@
+/* Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ * Copyright: Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.reactive.context.impl;
+
+import java.util.concurrent.ConcurrentMap;
+
+import io.vertx.core.impl.VertxBuilder;
+import io.vertx.core.spi.VertxServiceProvider;
+import io.vertx.core.spi.context.storage.ContextLocal;
+
+/**
+ * SPI Implementation for {@link ContextLocal} storage.
+ */
+public class ContextualDataStorage implements VertxServiceProvider {
+
+	@SuppressWarnings("rawtypes")
+	static ContextLocal<ConcurrentMap> CONTEXTUAL_DATA_KEY = ContextLocal.registerLocal( ConcurrentMap.class );
+
+	@Override
+	public void init(VertxBuilder vertxBuilder) {
+	}
+}

hibernate-reactive-core/src/main/java/org/hibernate/reactive/context/impl/VertxContext.java

@@ -6,9 +6,12 @@
 package org.hibernate.reactive.context.impl;
 
 import java.lang.invoke.MethodHandles;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
 
 import io.vertx.core.Vertx;
 import io.vertx.core.impl.ContextInternal;
+import io.vertx.core.spi.context.storage.AccessMode;
 
 import org.hibernate.reactive.context.Context;
 import org.hibernate.reactive.logging.impl.Log;
@@ -36,10 +39,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 = ContextInternal.current();
 		if ( context != null ) {
 			if ( trace ) LOG.tracef( "Putting key,value in context: [%1$s, %2$s]", key, instance );
-			context.putLocal( key, instance );
+			VertxContext.<T>contextualDataMap( context ).put( key, instance );
 		}
 		else {
 			if ( trace ) LOG.tracef( "Context is null for key,value: [%1$s, %2$s]", key, instance  );
@@ -49,9 +52,9 @@ public <T> void put(Key<T> key, T instance) {
 
 	@Override
 	public <T> T get(Key<T> key) {
-		final io.vertx.core.Context context = Vertx.currentContext();
+		final ContextInternal context = ContextInternal.current();
 		if ( context != null ) {
-			T local = context.getLocal( key );
+			T local = VertxContext.<T>contextualDataMap( context ).get( key );
 			if ( trace ) LOG.tracef( "Getting value %2$s from context for key %1$s", key, local  );
 			return local;
 		}
@@ -63,9 +66,9 @@ public <T> T get(Key<T> key) {
 
 	@Override
 	public void remove(Key<?> key) {
-		final io.vertx.core.Context context = Vertx.currentContext();
+		final ContextInternal context = ContextInternal.current();
 		if ( context != null ) {
-			boolean removed = context.removeLocal( key );
+			boolean removed = contextualDataMap( context ).remove( key ) != null;
 			if ( trace ) LOG.tracef( "Key %s removed from context: %s", key, removed );
 		}
 		else {
@@ -92,4 +95,12 @@ public void execute(Runnable runnable) {
 		}
 	}
 
+	@SuppressWarnings({ "unchecked" })
+	private static <T> ConcurrentMap<Key<T>, T> contextualDataMap(ContextInternal vertxContext) {
+		return vertxContext.getLocal(
+				ContextualDataStorage.CONTEXTUAL_DATA_KEY,
+				AccessMode.CONCURRENT,
+				ConcurrentHashMap::new
+		);
+	}
 }

hibernate-reactive-core/src/main/resources/META-INF/services/io.vertx.core.spi.VertxServiceProvider

@@ -0,0 +1 @@
+org.hibernate.reactive.context.impl.ContextualDataStorage