Proto2Coder: recompute the extension registry when mutated

Proto2Coder's deprecated methods mutate the existing coder's
extension host set, but it may already have been memoized.

----Release Notes----
[]
-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=115407399


Project: http://git-wip-us.apache.org/repos/asf/incubator-beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/incubator-beam/commit/6926d8ee
Tree: http://git-wip-us.apache.org/repos/asf/incubator-beam/tree/6926d8ee
Diff: http://git-wip-us.apache.org/repos/asf/incubator-beam/diff/6926d8ee

Branch: refs/heads/master
Commit: 6926d8ee1d5fa2afece0024e64e18dc06ec857c4
Parents: 96b02f4
Author: dhalperi <[email protected]>
Authored: Tue Feb 23 18:38:04 2016 -0800
Committer: Davor Bonaci <[email protected]>
Committed: Thu Feb 25 23:58:27 2016 -0800

----------------------------------------------------------------------
 .../google/cloud/dataflow/sdk/coders/Proto2Coder.java   | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/6926d8ee/sdk/src/main/java/com/google/cloud/dataflow/sdk/coders/Proto2Coder.java
----------------------------------------------------------------------
diff --git 
a/sdk/src/main/java/com/google/cloud/dataflow/sdk/coders/Proto2Coder.java 
b/sdk/src/main/java/com/google/cloud/dataflow/sdk/coders/Proto2Coder.java
index b107f37..56ec0c7 100644
--- a/sdk/src/main/java/com/google/cloud/dataflow/sdk/coders/Proto2Coder.java
+++ b/sdk/src/main/java/com/google/cloud/dataflow/sdk/coders/Proto2Coder.java
@@ -198,6 +198,11 @@ public class Proto2Coder<T extends Message> extends 
AtomicCoder<T> {
       }
       extensionHostClasses.add(extensionHost);
     }
+    // The memoized extension registry needs to be recomputed because we have 
mutated this object.
+    synchronized (this) {
+      memoizedExtensionRegistry = null;
+      getExtensionRegistry();
+    }
     return this;
   }
 
@@ -293,18 +298,19 @@ public class Proto2Coder<T extends Message> extends 
AtomicCoder<T> {
 
   private transient ExtensionRegistry memoizedExtensionRegistry;
 
-  private ExtensionRegistry getExtensionRegistry() {
+  private synchronized ExtensionRegistry getExtensionRegistry() {
     if (memoizedExtensionRegistry == null) {
-      memoizedExtensionRegistry = ExtensionRegistry.newInstance();
+      ExtensionRegistry registry = ExtensionRegistry.newInstance();
       for (Class<?> extensionHost : extensionHostClasses) {
         try {
           extensionHost
               .getDeclaredMethod("registerAllExtensions", 
ExtensionRegistry.class)
-              .invoke(null, memoizedExtensionRegistry);
+              .invoke(null, registry);
         } catch (IllegalAccessException | InvocationTargetException | 
NoSuchMethodException e) {
           throw new IllegalStateException(e);
         }
       }
+      memoizedExtensionRegistry = registry.getUnmodifiable();
     }
     return memoizedExtensionRegistry;
   }

Reply via email to