Update Precondition.checkNotNull to Precondition.checkArgument

Author: zhouyejoe

common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/RemoteBlockPushResolver.java

@@ -87,7 +87,7 @@ public class RemoteBlockPushResolver implements MergedShuffleFileManager {
    */
   private final ConcurrentMap<String, AppShuffleInfo> appsShuffleInfo;
 
-  private final Executor mergedShuffleCleanerExecutor;
+  private final Executor mergedShuffleCleaner;
   private final TransportConf conf;
   private final int minChunkSize;
   private final int ioExceptionsThresholdDuringMerge;
@@ -100,7 +100,7 @@ public class RemoteBlockPushResolver implements MergedShuffleFileManager {
   public RemoteBlockPushResolver(TransportConf conf) {
     this.conf = conf;
     this.appsShuffleInfo = new ConcurrentHashMap<>();
-    this.mergedShuffleCleanerExecutor = Executors.newSingleThreadExecutor(
+    this.mergedShuffleCleaner = Executors.newSingleThreadExecutor(
       // Add `spark` prefix because it will run in NM in Yarn mode.
       NettyUtils.createThreadFactory("spark-shuffle-merged-shuffle-directory-cleaner"));
     this.minChunkSize = conf.minChunkSizeInMergedShuffleFile();
@@ -121,9 +121,9 @@ public ShuffleIndexInformation load(File file) throws IOException {
   @VisibleForTesting
   protected AppShuffleInfo validateAndGetAppShuffleInfo(String appId) {
     // TODO: [SPARK-33236] Change the message when this service is able to handle NM restart
-    AppShuffleInfo appShuffleInfo =
-      Preconditions.checkNotNull(appsShuffleInfo.get(appId),
-        "application " + appId + " is not registered or NM was restarted.");
+    AppShuffleInfo appShuffleInfo = appsShuffleInfo.get(appId);
+    Preconditions.checkArgument(appShuffleInfo != null,
+      "application " + appId + " is not registered or NM was restarted.");
     return appShuffleInfo;
   }
 
@@ -255,7 +255,7 @@ public void applicationRemoved(String appId, boolean cleanupLocalDirs) {
     logger.info("Application {} removed, cleanupLocalDirs = {}", appId, cleanupLocalDirs);
     AppShuffleInfo appShuffleInfo = appsShuffleInfo.remove(appId);
     if (null != appShuffleInfo) {
-      mergedShuffleCleanerExecutor.execute(
+      mergedShuffleCleaner.execute(
         () -> closeAndDeletePartitionFilesIfNeeded(appShuffleInfo, cleanupLocalDirs));
     }
   }
@@ -501,7 +501,7 @@ public void registerExecutor(String appId, ExecutorShuffleInfo executorInfo) {
             AppShuffleInfo appShuffleInfo = originalAppShuffleInfo.get();
             logger.warn("Cleanup shuffle info and merged shuffle files for {}_{} as new " +
                 "application attempt registered", appId, appShuffleInfo.attemptId);
-            mergedShuffleCleanerExecutor.execute(
+            mergedShuffleCleaner.execute(
               () -> closeAndDeletePartitionFilesIfNeeded(appShuffleInfo, true));
           }
         }
@@ -538,10 +538,13 @@ private PushBlockStreamCallback(
         String streamId,
         AppShufflePartitionInfo partitionInfo,
         int mapIndex) {
-      this.mergeManager = Preconditions.checkNotNull(mergeManager);
-      this.appShuffleInfo = Preconditions.checkNotNull(appShuffleInfo);
+      Preconditions.checkArgument(mergeManager != null);
+      this.mergeManager = mergeManager;
+      Preconditions.checkArgument(appShuffleInfo != null);
+      this.appShuffleInfo = appShuffleInfo;
       this.streamId = streamId;
-      this.partitionInfo = Preconditions.checkNotNull(partitionInfo);
+      Preconditions.checkArgument(partitionInfo != null);
+      this.partitionInfo = partitionInfo;
       this.mapIndex = mapIndex;
       abortIfNecessary();
     }
@@ -855,7 +858,8 @@ public static class AppShufflePartitionInfo {
         File dataFile,
         MergeShuffleFile indexFile,
         MergeShuffleFile metaFile) throws IOException {
-      this.appId = Preconditions.checkNotNull(appId, "app id is null");
+      Preconditions.checkArgument(appId != null, "app id is null");
+      this.appId = appId;
       this.shuffleId = shuffleId;
       this.reduceId = reduceId;
       this.dataChannel = new FileOutputStream(dataFile).getChannel();

common/network-shuffle/src/test/java/org/apache/spark/network/shuffle/RemoteBlockPushResolverSuite.java

@@ -407,7 +407,7 @@ public void testFailureInAStreamDoesNotInterfereWithStreamWhichIsWriting() throw
     }
   }
 
-  @Test(expected = NullPointerException.class)
+  @Test(expected = IllegalArgumentException.class)
   public void testUpdateLocalDirsOnlyOnce() throws IOException {
     String testApp = "updateLocalDirsOnlyOnceTest";
     Path[] activeLocalDirs = createLocalDirs(1);
@@ -427,10 +427,10 @@ public void testUpdateLocalDirsOnlyOnce() throws IOException {
     removeApplication(testApp);
     try {
       pushResolver.getMergedBlockDirs(testApp);
-    } catch (Throwable e) {
-      assertTrue(e.getMessage()
-        .startsWith("application " + testApp + " is not registered or NM was restarted."));
-      Throwables.propagate(e);
+    } catch (IllegalArgumentException e) {
+      assertEquals(e.getMessage(),
+        "application " + testApp + " is not registered or NM was restarted.");
+      throw e;
     }
   }
 
@@ -450,7 +450,7 @@ public void testExecutorRegisterWithInvalidJsonForPushShuffle() throws IOExcepti
     }
   }
 
-  @Test(expected = NullPointerException.class)
+  @Test(expected = IllegalArgumentException.class)
   public void testExecutorRegistrationFromTwoAppAttempts() throws IOException {
     String testApp = "testExecutorRegistrationFromTwoAppAttempts";
     Path[] attempt1LocalDirs = createLocalDirs(1);
@@ -480,10 +480,10 @@ public void testExecutorRegistrationFromTwoAppAttempts() throws IOException {
     removeApplication(testApp);
     try {
       pushResolver.getMergedBlockDirs(testApp);
-    } catch (Throwable e) {
-      assertTrue(e.getMessage()
-        .startsWith("application " + testApp + " is not registered or NM was restarted."));
-      Throwables.propagate(e);
+    } catch (IllegalArgumentException e) {
+      assertEquals(e.getMessage(),
+        "application " + testApp + " is not registered or NM was restarted.");
+      throw e;
     }
   }