From 87390cb60aa970a1ffa1dd39e56c060176ab1b4b Mon Sep 17 00:00:00 2001
From: Igor Melnichenko <i-melnichenko@yandex-team.ru>
Date: Tue, 4 Aug 2026 23:12:52 +0300
Subject: [PATCH] Making CodecRegistry safe for concurrent use

One CodecRegistry instance is shared by every reader and writer created
from a TopicClient. getCodec is called from the compression threads
(WriterQueue), from the decompression threads (MessageDecoder via
Encoder.decode) and from gRPC callbacks, while registerCodec is public
API that may be called at any time. Backing that with a plain HashMap
means concurrent puts can lose entries or corrupt the table.

Use a ConcurrentHashMap. The new test loses codecs on every run with the
old implementation.
---
 .../ydb/topic/description/CodecRegistry.java  |  7 +-
 .../ydb/topic/impl/CodecRegistryTest.java     | 77 ++++++++++++++++++-
 2 files changed, 77 insertions(+), 7 deletions(-)

diff --git a/topic/src/main/java/tech/ydb/topic/description/CodecRegistry.java b/topic/src/main/java/tech/ydb/topic/description/CodecRegistry.java
index 513201770..a09dd8e54 100644
--- a/topic/src/main/java/tech/ydb/topic/description/CodecRegistry.java
+++ b/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);
         }
diff --git a/topic/src/test/java/tech/ydb/topic/impl/CodecRegistryTest.java b/topic/src/test/java/tech/ydb/topic/impl/CodecRegistryTest.java
index 35e0baa77..5841da5e3 100644
--- a/topic/src/test/java/tech/ydb/topic/impl/CodecRegistryTest.java
+++ b/topic/src/test/java/tech/ydb/topic/impl/CodecRegistryTest.java
@@ -1,5 +1,16 @@
 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;
@@ -7,10 +18,6 @@
 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);