From 87390cb60aa970a1ffa1dd39e56c060176ab1b4b Mon Sep 17 00:00:00 2001 From: Igor Melnichenko Date: Tue, 4 Aug 2026 23:12:52 +0300 Subject: [PATCH 1/2] 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 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 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> 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); From a4a67f10a71ba86f2e327e749e02c84898bd75e9 Mon Sep 17 00:00:00 2001 From: Alexandr Gorshenin Date: Wed, 12 Aug 2026 15:43:45 +0100 Subject: [PATCH 2/2] Updated CodecRegistryTest --- .../ydb/topic/description/CodecRegistry.java | 5 +- .../ydb/topic/impl/CodecRegistryTest.java | 87 +++---------------- 2 files changed, 14 insertions(+), 78 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 a09dd8e54..04df98b78 100644 --- a/topic/src/main/java/tech/ydb/topic/description/CodecRegistry.java +++ b/topic/src/main/java/tech/ydb/topic/description/CodecRegistry.java @@ -33,11 +33,12 @@ public CodecRegistry() { * @return previous implementation with associated codec */ public Codec registerCodec(Codec codec) { - assert codec != null; + if (codec == null) { + throw new IllegalArgumentException("Codec must be not null"); + } int codecId = codec.getId(); Codec result = customCodecMap.put(codecId, codec); - if (result != null) { logger.info( "Replace codec which have already associated with this id. CodecId: {} 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 5841da5e3..b341c1093 100644 --- a/topic/src/test/java/tech/ydb/topic/impl/CodecRegistryTest.java +++ b/topic/src/test/java/tech/ydb/topic/impl/CodecRegistryTest.java @@ -3,13 +3,6 @@ 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; @@ -26,7 +19,7 @@ public class CodecRegistryTest { CodecRegistry registry; - private static final int codecId = 10113; + private static final int CUSTOM_ID = 10113; @Before public void beforeTest() { @@ -41,15 +34,14 @@ public void registerCustomCodecShouldDoubleRegisterCodecAndReturnLastCodec() { registry.registerCodec(codec1); Assert.assertEquals(codec1, registry.registerCodec(codec2)); - Assert.assertEquals(codec2, registry.getCodec(codecId)); - Assert.assertNotEquals(codec1, registry.getCodec(codecId)); + Assert.assertEquals(codec2, registry.getCodec(CUSTOM_ID)); + Assert.assertNotEquals(codec1, registry.getCodec(CUSTOM_ID)); } @Test public void registerCustomCodecShouldNotAcceptNull() { - Assert.assertThrows( - AssertionError.class, - () -> registry.registerCodec(null)); + Exception ex = Assert.assertThrows(IllegalArgumentException.class, () -> registry.registerCodec(null)); + Assert.assertEquals("Codec must be not null", ex.getMessage()); } @Test @@ -61,68 +53,6 @@ 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> 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); @@ -134,7 +64,12 @@ static class CodecTopic implements Codec { int codec; public CodecTopic() { - this.codec = codecId; + this.codec = CUSTOM_ID; + } + + @Override + public String toString() { + return "CustomCodec"; } public void setCodecId(int codecId) {