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);