YARN-11463. Node Labels root directory creation doesn't have a retry logic - addendum (#5614)

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/nodelabels/store/AbstractFSNodeStore.java

@@ -69,9 +69,9 @@ public abstract class AbstractFSNodeStore<M> {
     int maxRetries = conf.getInt(YarnConfiguration.NODE_STORE_ROOT_DIR_NUM_RETRIES,
         YarnConfiguration.NODE_STORE_ROOT_DIR_NUM_DEFAULT_RETRIES);
     int retryCount = 0;
-    boolean success = fs.mkdirs(fsWorkingPath);
+    boolean success = false;
 
-    while (!success && retryCount < maxRetries) {
+    while (!success && retryCount <= maxRetries) {
       try {
         if (!fs.exists(fsWorkingPath)) {
           success = fs.mkdirs(fsWorkingPath);
@@ -80,7 +80,7 @@ public abstract class AbstractFSNodeStore<M> {
         }
       } catch (IOException e) {
         retryCount++;
-        if (retryCount >= maxRetries) {
+        if (retryCount > maxRetries) {
           throw e;
         }
         try {

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/test/java/org/apache/hadoop/yarn/nodelabels/TestFileSystemNodeLabelsStore.java

@@ -348,24 +348,56 @@ public class TestFileSystemNodeLabelsStore extends NodeLabelTestBase {
 
   @MethodSource("getParameters")
   @ParameterizedTest
-  void testRootMkdirOnInitStore(String className) throws Exception {
+  void testRootMkdirOnInitStoreWhenRootDirectoryAlreadyExists(String className) throws Exception {
     initTestFileSystemNodeLabelsStore(className);
     final FileSystem mockFs = Mockito.mock(FileSystem.class);
+    final FileSystemNodeLabelsStore mockStore = createMockNodeLabelsStore(mockFs);
+    final int expectedMkdirsCount = 0;
+
+    Mockito.when(mockStore.getFs().exists(Mockito.any(Path.class)))
+        .thenReturn(true);
+    verifyMkdirsCount(mockStore, expectedMkdirsCount);
+  }
+
+  @MethodSource("getParameters")
+  @ParameterizedTest
+  void testRootMkdirOnInitStoreWhenRootDirectoryNotExists(String className) throws Exception {
+    initTestFileSystemNodeLabelsStore(className);
+    final FileSystem mockFs = Mockito.mock(FileSystem.class);
+    final FileSystemNodeLabelsStore mockStore = createMockNodeLabelsStore(mockFs);
+    final int expectedMkdirsCount = 1;
+
+    Mockito.when(mockStore.getFs().exists(Mockito.any(Path.class)))
+        .thenReturn(false).thenReturn(true);
+    verifyMkdirsCount(mockStore, expectedMkdirsCount);
+  }
+
+  @MethodSource("getParameters")
+  @ParameterizedTest
+  void testRootMkdirOnInitStoreRetryLogic(String className) throws Exception {
+    initTestFileSystemNodeLabelsStore(className);
+    final FileSystem mockFs = Mockito.mock(FileSystem.class);
+    final FileSystemNodeLabelsStore mockStore = createMockNodeLabelsStore(mockFs);
+    final int expectedMkdirsCount = 2;
+
+    Mockito.when(mockStore.getFs().exists(Mockito.any(Path.class)))
+        .thenReturn(false).thenReturn(false).thenReturn(true);
+    verifyMkdirsCount(mockStore, expectedMkdirsCount);
+  }
+
+  private FileSystemNodeLabelsStore createMockNodeLabelsStore(FileSystem mockFs) {
     FileSystemNodeLabelsStore mockStore = new FileSystemNodeLabelsStore() {
       public void initFileSystem(Configuration config) throws IOException {
         setFs(mockFs);
       }
     };
-
     mockStore.setFs(mockFs);
-    verifyMkdirsCount(mockStore, true, 1);
+    return mockStore;
   }
 
   private void verifyMkdirsCount(FileSystemNodeLabelsStore store,
-      boolean existsRetVal, int expectedNumOfCalls)
+      int expectedNumOfCalls)
       throws Exception {
-    Mockito.when(store.getFs().exists(Mockito.any(
-        Path.class))).thenReturn(existsRetVal);
     store.init(conf, mgr);
     Mockito.verify(store.getFs(), Mockito.times(
         expectedNumOfCalls)).mkdirs(Mockito.any(Path