[GitHub] [incubator-druid] jihoonson commented on a change in pull request #8114: Fix race between canHandle() and addSegment() in StorageLocation

[GitBox <gi...@apache.org>]

jihoonson commented on a change in pull request #8114: Fix race between canHandle() and addSegment() in StorageLocation
URL: https://github.com/apache/incubator-druid/pull/8114#discussion_r306573832
 
 

 ##########
 File path: server/src/main/java/org/apache/druid/segment/loading/StorageLocation.java
 ##########
 @@ -72,65 +83,81 @@ public long getMaxSize()
   }
 
   /**
-   * Add a new file to this location. The given file argument must be a file rather than directory.
+   * Remove a segment file from this location. The given file argument must be a file rather than directory.
    */
-  public synchronized void addFile(File file)
+  public synchronized void removeFile(File file)
   {
-    if (file.isDirectory()) {
-      throw new ISE("[%s] must be a file. Use a");
-    }
-    if (files.add(file)) {
-      currSize += FileUtils.sizeOf(file);
+    if (files.remove(file)) {
+      currSize -= FileUtils.sizeOf(file);
+    } else {
+      log.warn("File[%s] is not found under this location[%s]", file, path);
     }
   }
 
   /**
-   * Add a new segment dir to this location. The segment size is added to currSize.
+   * Remove a segment dir from this location. The segment size is subtracted from currSize.
    */
-  public synchronized void addSegmentDir(File segmentDir, DataSegment segment)
+  public synchronized void removeSegmentDir(File segmentDir, DataSegment segment)
   {
-    if (files.add(segmentDir)) {
-      currSize += segment.getSize();
+    if (files.remove(segmentDir)) {
+      currSize -= segment.getSize();
+    } else {
+      log.warn("SegmentDir[%s] is not found under this location[%s]", segmentDir, path);
     }
   }
 
   /**
-   * Remove a segment file from this location. The given file argument must be a file rather than directory.
+   * Reserves space to store the given segment. The segment size is added to currSize.
+   * If it succeeds, it returns a file for the given segmentDir in this storage location. Returns null otherwise.
    */
-  public synchronized void removeFile(File file)
+  @Nullable
+  public synchronized File reserve(String segmentDir, DataSegment segment)
   {
-    if (files.remove(file)) {
-      currSize -= FileUtils.sizeOf(file);
-    }
+    return reserve(segmentDir, segment.getId(), segment.getSize());
   }
 
   /**
-   * Remove a segment dir from this location. The segment size is subtracted from currSize.
+   * Reserves space to store the given segment.
+   * If it succeeds, it returns a file for the given segmentFilePathToAdd in this storage location.
+   * Returns null otherwise.
    */
-  public synchronized void removeSegmentDir(File segmentDir, DataSegment segment)
+  @Nullable
+  public synchronized File reserve(String segmentFilePathToAdd, SegmentId segmentId, long segmentSize)
   {
-    if (files.remove(segmentDir)) {
-      currSize -= segment.getSize();
+    final File segmentFileToAdd = new File(path, segmentFilePathToAdd);
+    if (files.contains(segmentFileToAdd)) {
+      return null;
+    }
+    if (canHandle(segmentId, segmentSize)) {
+      files.add(segmentFileToAdd);
+      currSize += segmentSize;
+      return segmentFileToAdd;
+    } else {
+      return null;
     }
   }
 
-  public boolean canHandle(DataSegment segment)
+  /**
+   * This method is available for only unit tests. Production code must use {@link #reserve} instead.
 
 Review comment:
   Tried to make it clear.

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services

---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org