From b233609def085954fa81cd72506491d32b5351df Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dominik=20D=C4=99bowczyk?= <dominikd@spotify.com>
Date: Thu, 21 Sep 2023 17:21:58 +0200
Subject: [PATCH] Make SerializableConfiguration cacheable

---
 .../sdk/io/hadoop/SerializableConfiguration.java | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/sdks/java/io/hadoop-common/src/main/java/org/apache/beam/sdk/io/hadoop/SerializableConfiguration.java b/sdks/java/io/hadoop-common/src/main/java/org/apache/beam/sdk/io/hadoop/SerializableConfiguration.java
index bb59b07cc280c..7311035e69d2c 100644
--- a/sdks/java/io/hadoop-common/src/main/java/org/apache/beam/sdk/io/hadoop/SerializableConfiguration.java
+++ b/sdks/java/io/hadoop-common/src/main/java/org/apache/beam/sdk/io/hadoop/SerializableConfiguration.java
@@ -17,6 +17,8 @@
  */
 package org.apache.beam.sdk.io.hadoop;
 
+import java.io.ByteArrayOutputStream;
+import java.io.DataOutputStream;
 import java.io.Externalizable;
 import java.io.IOException;
 import java.io.ObjectInput;
@@ -39,6 +41,8 @@ public class SerializableConfiguration implements Externalizable {
 
   private transient Configuration conf;
 
+  private transient byte[] serializationCache;
+
   public SerializableConfiguration() {}
 
   public SerializableConfiguration(Configuration conf) {
@@ -49,13 +53,23 @@ public SerializableConfiguration(Configuration conf) {
   }
 
   public Configuration get() {
+    if(serializationCache != null) {
+      serializationCache = null;
+    }
     return conf;
   }
 
   @Override
   public void writeExternal(ObjectOutput out) throws IOException {
+    if (serializationCache == null) {
+      ByteArrayOutputStream baos = new ByteArrayOutputStream(512);
+      try (DataOutputStream dos = new DataOutputStream(baos)) {
+        conf.write(dos);
+        serializationCache = baos.toByteArray();
+      }
+    }
     out.writeUTF(conf.getClass().getCanonicalName());
-    conf.write(out);
+    out.write(serializationCache);
   }
 
   @Override