Add synchronized to getExtensionRegistry() lazy initialization

jakubrauch · Apr 27, 2023 · 36b5203 · 36b5203
1 parent 3c80555
commit 36b5203
Showing 1 changed file with 14 additions and 10 deletions.
diff --git a/...extensions/protobuf/src/main/java/org/apache/beam/sdk/extensions/protobuf/ProtoCoder.java b/...extensions/protobuf/src/main/java/org/apache/beam/sdk/extensions/protobuf/ProtoCoder.java
@@ -247,17 +247,21 @@ public Set<Class<?>> getExtensionHosts() {
    */
   public ExtensionRegistry getExtensionRegistry() {
     if (memoizedExtensionRegistry == null) {
-      ExtensionRegistry registry = ExtensionRegistry.newInstance();
-      for (Class<?> extensionHost : extensionHostClasses) {
-        try {
-          extensionHost
-              .getDeclaredMethod("registerAllExtensions", ExtensionRegistry.class)
-              .invoke(null, registry);
-        } catch (IllegalAccessException | InvocationTargetException | NoSuchMethodException e) {
-          throw new IllegalStateException(e);
+      synchronized (this) {
+        if (memoizedExtensionRegistry == null) {
+          ExtensionRegistry registry = ExtensionRegistry.newInstance();
+          for (Class<?> extensionHost : extensionHostClasses) {
+            try {
+              extensionHost
+                  .getDeclaredMethod("registerAllExtensions", ExtensionRegistry.class)
+                  .invoke(null, registry);
+            } catch (IllegalAccessException | InvocationTargetException | NoSuchMethodException e) {
+              throw new IllegalStateException(e);
+            }
+          }
+          memoizedExtensionRegistry = registry.getUnmodifiable();
         }
       }
-      memoizedExtensionRegistry = registry.getUnmodifiable();
     }
     return memoizedExtensionRegistry;
   }
@@ -275,7 +279,7 @@ public ExtensionRegistry getExtensionRegistry() {
   final Set<Class<?>> extensionHostClasses;
 
   // Transient fields that are lazy initialized and then memoized.
-  private transient ExtensionRegistry memoizedExtensionRegistry;
+  private transient volatile ExtensionRegistry memoizedExtensionRegistry;
   transient Parser<T> memoizedParser;
 
   /** Private constructor. */