6363import java .util .concurrent .atomic .AtomicReference ;
6464import java .util .concurrent .locks .Condition ;
6565import java .util .concurrent .locks .Lock ;
66+ import java .util .concurrent .locks .ReadWriteLock ;
6667import java .util .concurrent .locks .ReentrantLock ;
68+ import java .util .concurrent .locks .ReentrantReadWriteLock ;
6769
6870import static org .apache .doris .flink .sink .LoadStatus .PUBLISH_TIMEOUT ;
6971import static org .apache .doris .flink .sink .LoadStatus .SUCCESS ;
@@ -111,6 +113,7 @@ public class DorisBatchStreamLoad implements Serializable {
111113 private final AtomicLong currentCacheBytes = new AtomicLong (0L );
112114 private final Lock lock = new ReentrantLock ();
113115 private final Condition block = lock .newCondition ();
116+ private final Map <String , ReadWriteLock > bufferMapLock = new ConcurrentHashMap <>();
114117
115118 public DorisBatchStreamLoad (
116119 DorisOptions dorisOptions ,
@@ -181,7 +184,7 @@ public DorisBatchStreamLoad(
181184 public void writeRecord (String database , String table , byte [] record ) {
182185 checkFlushException ();
183186 String bufferKey = getTableIdentifier (database , table );
184-
187+ getLock ( bufferKey ). readLock (). lock ();
185188 BatchRecordBuffer buffer =
186189 bufferMap .computeIfAbsent (
187190 bufferKey ,
@@ -194,6 +197,7 @@ public void writeRecord(String database, String table, byte[] record) {
194197
195198 int bytes = buffer .insert (record );
196199 currentCacheBytes .addAndGet (bytes );
200+ getLock (bufferKey ).readLock ().unlock ();
197201 if (currentCacheBytes .get () > maxBlockedBytes ) {
198202 lock .lock ();
199203 try {
@@ -283,11 +287,19 @@ private synchronized boolean flush(String bufferKey, boolean waitUtilDone) {
283287 }
284288
285289 private synchronized void flushBuffer (String bufferKey ) {
286- BatchRecordBuffer buffer = bufferMap .get (bufferKey );
290+ BatchRecordBuffer buffer ;
291+ try {
292+ getLock (bufferKey ).writeLock ().lock ();
293+ buffer = bufferMap .remove (bufferKey );
294+ } finally {
295+ getLock (bufferKey ).writeLock ().unlock ();
296+ }
297+ if (buffer == null ) {
298+ return ;
299+ }
287300 buffer .setLabelName (labelGenerator .generateBatchLabel (buffer .getTable ()));
288301 LOG .debug ("flush buffer for key {} with label {}" , bufferKey , buffer .getLabelName ());
289302 putRecordToFlushQueue (buffer );
290- bufferMap .remove (bufferKey );
291303 }
292304
293305 private void putRecordToFlushQueue (BatchRecordBuffer buffer ) {
@@ -374,6 +386,10 @@ private boolean merge(BatchRecordBuffer mergeBuffer, BatchRecordBuffer buffer) {
374386 return true ;
375387 }
376388
389+ private ReadWriteLock getLock (String bufferKey ) {
390+ return bufferMapLock .computeIfAbsent (bufferKey , k -> new ReentrantReadWriteLock ());
391+ }
392+
377393 class LoadAsyncExecutor implements Runnable {
378394
379395 private int flushQueueSize ;
0 commit comments