Making CodecRegistry thread-safe

topic/src/main/java/tech/ydb/topic/description/CodecRegistry.java

@@ -1,7 +1,7 @@
 package tech.ydb.topic.description;
 
-import java.util.HashMap;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -17,10 +17,11 @@ public class CodecRegistry {
 
     private static final Logger logger = LoggerFactory.getLogger(CodecRegistry.class);
 
-    final Map<Integer, Codec> customCodecMap;
+    // registerCodec may be called at any moment, while getCodec is used by the compression and the
+    // decompression threads of every reader and writer created from the same TopicClient
+    private final Map<Integer, Codec> customCodecMap = new ConcurrentHashMap<>();
 
     public CodecRegistry() {
-        customCodecMap = new HashMap<>();
         for (Codec codec: StandardCodecs.getAvailableCodecs()) {
             customCodecMap.put(codec.getId(), codec);
         }

topic/src/test/java/tech/ydb/topic/impl/CodecRegistryTest.java

@@ -1,16 +1,23 @@
 package tech.ydb.topic.impl;
 
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.atomic.AtomicBoolean;
+
 import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
 
 import tech.ydb.topic.description.Codec;
 import tech.ydb.topic.description.CodecRegistry;
 
-import java.io.IOException;
-import java.io.InputStream;
-import java.io.OutputStream;
-
 /**
  * Unit tests for check simple logic for register custom codec
  *
@@ -54,6 +61,68 @@ public void registerCustomCodecShouldRegisterAndOverrideAnyCodec() {
         expectRegisterCodec(4, codec1, ZstdCodec.getInstance());
     }
 
+    @Test(timeout = 60_000)
+    public void registerCustomCodecIsSafeForConcurrentUse() throws Exception {
+        int writerCount = 4;
+        int codecsPerWriter = 500;
+        int firstCodecId = 20000;
+
+        ExecutorService executor = Executors.newFixedThreadPool(writerCount + 1);
+
+        try {
+            CountDownLatch start = new CountDownLatch(1);
+            AtomicBoolean readersRun = new AtomicBoolean(true);
+            List<Future<?>> futures = new ArrayList<>();
+
+            for (int writer = 0; writer < writerCount; writer += 1) {
+                int base = firstCodecId + writer * codecsPerWriter;
+                futures.add(executor.submit(() -> {
+                    start.await();
+
+                    for (int idx = 0; idx < codecsPerWriter; idx += 1) {
+                        CodecTopic codec = new CodecTopic();
+                        codec.setCodecId(base + idx);
+                        registry.registerCodec(codec);
+                    }
+
+                    return null;
+                }));
+            }
+
+            futures.add(executor.submit(() -> {
+                start.await();
+
+                while (readersRun.get()) {
+                    registry.getCodec(Codec.RAW);
+                }
+
+                return null;
+            }));
+
+            start.countDown();
+
+            try {
+                for (int i = 0; i < writerCount; i += 1) {
+                    futures.get(i).get();
+                }
+
+                readersRun.set(false);
+                futures.get(writerCount).get();
+            } finally {
+                readersRun.set(false);
+            }
+        } finally {
+            executor.shutdownNow();
+        }
+
+        for (int codecId = firstCodecId; codecId < firstCodecId + writerCount * codecsPerWriter; codecId += 1) {
+            Assert.assertNotNull(
+                    "codec " + codecId + " was lost by a concurrent registration",
+                    registry.getCodec(codecId)
+            );
+        }
+    }
+
     void expectRegisterCodec(int codecId, CodecTopic newCodec, Codec oldCodec) {
         newCodec.setCodecId(codecId);
         Codec codecOldPredefined = registry.registerCodec(newCodec);