extract public user context method to BaseFury

apache · May 7, 2024 · 74fa5f7 · 74fa5f7
1 parent 757e5fd
commit 74fa5f7
Show file tree

Hide file tree

Showing 5 changed files with 29 additions and 2 deletions.
diff --git a/java/fury-core/src/main/java/org/apache/fury/AbstractThreadSafeFury.java b/java/fury-core/src/main/java/org/apache/fury/AbstractThreadSafeFury.java
@@ -20,6 +20,8 @@
 package org.apache.fury;
 
 import java.util.function.Consumer;
+import java.util.function.Function;
+import org.apache.fury.resolver.UserContextResolver;
 import org.apache.fury.serializer.Serializer;
 import org.apache.fury.serializer.SerializerFactory;
 
@@ -59,5 +61,11 @@ public void setSerializerFactory(SerializerFactory serializerFactory) {
     processCallback(fury -> fury.setSerializerFactory(serializerFactory));
   }
 
+  @Override
+  public void registerUserContext(
+      String name, Function<Fury, UserContextResolver> userContextResolverCtr) {
+    processCallback(fury -> fury.registerUserContext(name, userContextResolverCtr));
+  }
+
   protected abstract void processCallback(Consumer<Fury> callback);
 }
diff --git a/java/fury-core/src/main/java/org/apache/fury/BaseFury.java b/java/fury-core/src/main/java/org/apache/fury/BaseFury.java
@@ -20,9 +20,11 @@
 package org.apache.fury;
 
 import java.io.OutputStream;
+import java.util.function.Function;
 import org.apache.fury.io.FuryInputStream;
 import org.apache.fury.io.FuryReadableChannel;
 import org.apache.fury.memory.MemoryBuffer;
+import org.apache.fury.resolver.UserContextResolver;
 import org.apache.fury.serializer.BufferCallback;
 import org.apache.fury.serializer.Serializer;
 import org.apache.fury.serializer.SerializerFactory;
@@ -73,6 +75,14 @@ public interface BaseFury {
 
   void setSerializerFactory(SerializerFactory serializerFactory);
 
+  /**
+   * Register user context.
+   *
+   * @param name user context name to register
+   * @param userContextResolverCtr constructor for usr context
+   */
+  void registerUserContext(String name, Function<Fury, UserContextResolver> userContextResolverCtr);
+
   /** Return serialized <code>obj</code> as a byte array. */
   byte[] serialize(Object obj);
 

diff --git a/java/fury-core/src/main/java/org/apache/fury/Fury.java b/java/fury-core/src/main/java/org/apache/fury/Fury.java
@@ -28,6 +28,7 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.function.Consumer;
+import java.util.function.Function;
 import javax.annotation.concurrent.NotThreadSafe;
 import org.apache.fury.builder.JITContext;
 import org.apache.fury.collection.ObjectArray;
@@ -52,6 +53,7 @@
 import org.apache.fury.resolver.NoRefResolver;
 import org.apache.fury.resolver.RefResolver;
 import org.apache.fury.resolver.SerializationContext;
+import org.apache.fury.resolver.UserContextResolver;
 import org.apache.fury.serializer.ArraySerializers;
 import org.apache.fury.serializer.BufferCallback;
 import org.apache.fury.serializer.BufferObject;
@@ -184,6 +186,12 @@ public void setSerializerFactory(SerializerFactory serializerFactory) {
     classResolver.setSerializerFactory(serializerFactory);
   }
 
+  @Override
+  public void registerUserContext(
+      String name, Function<Fury, UserContextResolver> userContextResolverCtr) {
+    serializationContext.registerUserContextResolver(name, userContextResolverCtr.apply(this));
+  }
+
   public SerializerFactory getSerializerFactory() {
     return classResolver.getSerializerFactory();
   }

diff --git a/java/fury-core/src/main/java/org/apache/fury/resolver/SerializationContext.java b/java/fury-core/src/main/java/org/apache/fury/resolver/SerializationContext.java
@@ -23,6 +23,7 @@
 import java.util.LinkedHashMap;
 import java.util.Map;
 import org.apache.fury.config.FuryBuilder;
+import org.apache.fury.util.Preconditions;
 
 /**
  * A context is used to add some context-related information, so that the serializers can set up
@@ -66,6 +67,7 @@ public Map<String, UserContextResolver> getUserContextResolvers() {
   }
 
   public void registerUserContextResolver(String name, UserContextResolver userContextResolver) {
+    Preconditions.checkState(!userContextResolvers.containsKey(name));
     userContextResolvers.put(name, userContextResolver);
   }
 

diff --git a/java/fury-core/src/test/java/org/apache/fury/resolver/UserContextTest.java b/java/fury-core/src/test/java/org/apache/fury/resolver/UserContextTest.java
@@ -111,8 +111,7 @@ public void checkShareUserContext() {
     List<StringUserContextResolver> resolvers = new ArrayList<>(NUM);
     for (int i = 0; i < NUM; i++) {
       final StringUserContextResolver userContextResolver = new StringUserContextResolver(fury);
-      fury.getSerializationContext()
-          .registerUserContextResolver(String.valueOf(i), userContextResolver);
+      fury.registerUserContext(String.valueOf(i), StringUserContextResolver::new);
       resolvers.add(userContextResolver);
     }
     final Foo o = Foo.create();