From a21f1f4c1e769333c9904109a8d468c61467b5ea Mon Sep 17 00:00:00 2001
From: "Doroszlai, Attila" <adoroszlai@apache.org>
Date: Thu, 9 Nov 2023 10:07:32 +0100
Subject: [PATCH] HDDS-9659. Intermittent failure in
 TestOzoneECClient#testPartialStripeWithPartialChunkRetry

---
 .../ozone/client/io/ECBlockOutputStreamEntry.java  |  5 +++--
 .../ozone/client/MockXceiverClientFactory.java     | 14 ++++++--------
 2 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/hadoop-ozone/client/src/main/java/org/apache/hadoop/ozone/client/io/ECBlockOutputStreamEntry.java b/hadoop-ozone/client/src/main/java/org/apache/hadoop/ozone/client/io/ECBlockOutputStreamEntry.java
index 8a8b32faaa7b..26d11f3d642b 100644
--- a/hadoop-ozone/client/src/main/java/org/apache/hadoop/ozone/client/io/ECBlockOutputStreamEntry.java
+++ b/hadoop-ozone/client/src/main/java/org/apache/hadoop/ozone/client/io/ECBlockOutputStreamEntry.java
@@ -91,16 +91,17 @@ public class ECBlockOutputStreamEntry extends BlockOutputStreamEntry {
   @Override
   void checkStream() throws IOException {
     if (!isInitialized()) {
-      blockOutputStreams =
+      final ECBlockOutputStream[] streams =
           new ECBlockOutputStream[replicationConfig.getRequiredNodes()];
       for (int i = currentStreamIdx; i < replicationConfig
           .getRequiredNodes(); i++) {
         List<DatanodeDetails> nodes = getPipeline().getNodes();
-        blockOutputStreams[i] =
+        streams[i] =
             new ECBlockOutputStream(getBlockID(), getXceiverClientManager(),
                 createSingleECBlockPipeline(getPipeline(), nodes.get(i), i + 1),
                 getBufferPool(), getConf(), getToken(), getClientMetrics());
       }
+      blockOutputStreams = streams;
     }
   }
 
diff --git a/hadoop-ozone/client/src/test/java/org/apache/hadoop/ozone/client/MockXceiverClientFactory.java b/hadoop-ozone/client/src/test/java/org/apache/hadoop/ozone/client/MockXceiverClientFactory.java
index 8c5fa762912b..02964da89746 100644
--- a/hadoop-ozone/client/src/test/java/org/apache/hadoop/ozone/client/MockXceiverClientFactory.java
+++ b/hadoop-ozone/client/src/test/java/org/apache/hadoop/ozone/client/MockXceiverClientFactory.java
@@ -25,12 +25,11 @@
 import java.io.IOException;
 import java.util.Collection;
 import java.util.Collections;
-import java.util.HashMap;
-import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 
 /**
  * Factory to create the mock datanode clients.
@@ -39,9 +38,9 @@ public class MockXceiverClientFactory
     implements XceiverClientFactory {
 
   private final Map<DatanodeDetails, MockDatanodeStorage> storage =
-      new HashMap<>();
+      new ConcurrentHashMap<>();
   private final Map<IOException, Set<DatanodeDetails>> pendingDNFailures =
-      new HashMap<>();
+      new ConcurrentHashMap<>();
 
   public void setFailedStorages(List<DatanodeDetails> failedStorages) {
     mockStorageFailure(failedStorages,
@@ -50,10 +49,9 @@ public void setFailedStorages(List<DatanodeDetails> failedStorages) {
 
   public void mockStorageFailure(Collection<DatanodeDetails> datanodes,
       IOException reason) {
-    if (!pendingDNFailures.containsKey(reason)) {
-      pendingDNFailures.put(reason, new HashSet<>());
-    }
-    pendingDNFailures.get(reason).addAll(datanodes);
+    pendingDNFailures
+        .computeIfAbsent(reason, k -> ConcurrentHashMap.newKeySet())
+        .addAll(datanodes);
     mockStorageFailure(reason);
   }