perf: Skip download of upgraded segments on historicals

processing/src/main/java/org/apache/druid/timeline/DataSegment.java

@@ -152,6 +152,14 @@ public static class PruneSpecsHolder
   @Nullable
   private final String indexingStateFingerprint;
 
+  /**
+   * ID of the segment that this segment was upgraded from during a concurrent APPEND + REPLACE operation.
+   * When non-null, the Historical can reuse the base segment's cached files instead of downloading from deep storage,
+   * since an upgraded segment shares the same underlying data files as its base segment.
+   */
+  @Nullable
+  private final SegmentId upgradedFromSegmentId;
+
   /**
    * @deprecated use {@link #builder(SegmentId)} or {@link #builder(DataSegment)} instead.
    */
@@ -183,6 +191,7 @@ public DataSegment(
         size,
         null,
         null,
+        null,
         PruneSpecsHolder.DEFAULT
     );
   }
@@ -219,6 +228,7 @@ public DataSegment(
         size,
         null,
         null,
+        null,
         PruneSpecsHolder.DEFAULT
     );
   }
@@ -259,6 +269,7 @@ public DataSegment(
         size,
         totalRows,
         indexingStateFingerprint,
+        null,
         pruneSpecsHolder
     );
   }
@@ -283,6 +294,7 @@ private DataSegment(
       @JsonProperty("size") long size,
       @JsonProperty("totalRows") Integer totalRows,
       @JsonProperty("indexingStateFingerprint") @Nullable String indexingStateFingerprint,
+      @JsonProperty("upgradedFromSegmentId") @Nullable SegmentId upgradedFromSegmentId,
       @JacksonInject PruneSpecsHolder pruneSpecsHolder
   )
   {
@@ -301,6 +313,7 @@ private DataSegment(
         size,
         totalRows,
         indexingStateFingerprint,
+        upgradedFromSegmentId,
         pruneSpecsHolder
     );
   }
@@ -320,6 +333,7 @@ public DataSegment(
       long size,
       Integer totalRows,
       String indexingStateFingerprint,
+      @Nullable SegmentId upgradedFromSegmentId,
       PruneSpecsHolder pruneSpecsHolder
   )
   {
@@ -343,6 +357,7 @@ public DataSegment(
     this.indexingStateFingerprint = indexingStateFingerprint == null ?
                                     null :
                                     STRING_INTERNER.intern(indexingStateFingerprint);
+    this.upgradedFromSegmentId = upgradedFromSegmentId;
   }
 
   /**
@@ -468,6 +483,19 @@ public String getIndexingStateFingerprint()
     return indexingStateFingerprint;
   }
 
+  @Nullable
+  @JsonProperty
+  @JsonInclude(JsonInclude.Include.NON_NULL)
+  public SegmentId getUpgradedFromSegmentId()
+  {
+    return upgradedFromSegmentId;
+  }
+
+  public DataSegment withUpgradedFromSegmentId(@Nullable SegmentId upgradedFromSegmentId)
+  {
+    return builder(this).upgradedFromSegmentId(upgradedFromSegmentId).build();
+  }
+
   @Override
   public boolean overshadows(DataSegment other)
   {
@@ -705,6 +733,8 @@ public static class Builder
     private long size;
     private Integer totalRows;
     private String indexingStateFingerprint;
+    @Nullable
+    private SegmentId upgradedFromSegmentId;
 
     /**
      * @deprecated use {@link #Builder(SegmentId)} or {@link #Builder(DataSegment)} instead.
@@ -751,6 +781,7 @@ private Builder(DataSegment segment)
       this.size = segment.getSize();
       this.totalRows = segment.getTotalRows();
       this.indexingStateFingerprint = segment.getIndexingStateFingerprint();
+      this.upgradedFromSegmentId = segment.getUpgradedFromSegmentId();
     }
 
     private Builder(DataSegment.Builder segmentBuilder)
@@ -769,6 +800,7 @@ private Builder(DataSegment.Builder segmentBuilder)
       this.size = segmentBuilder.size;
       this.totalRows = segmentBuilder.totalRows;
       this.indexingStateFingerprint = segmentBuilder.indexingStateFingerprint;
+      this.upgradedFromSegmentId = segmentBuilder.upgradedFromSegmentId;
     }
 
     public Builder dataSource(String dataSource)
@@ -855,6 +887,12 @@ public Builder indexingStateFingerprint(String indexingStateFingerprint)
       return this;
     }
 
+    public Builder upgradedFromSegmentId(@Nullable SegmentId upgradedFromSegmentId)
+    {
+      this.upgradedFromSegmentId = upgradedFromSegmentId;
+      return this;
+    }
+
     public DataSegment build()
     {
       // Check stuff that goes into the id, at least.
@@ -878,6 +916,7 @@ public DataSegment build()
           size,
           totalRows,
           indexingStateFingerprint,
+          upgradedFromSegmentId,
           PruneSpecsHolder.DEFAULT
       );
     }

server/src/main/java/org/apache/druid/segment/loading/SegmentLocalCacheManager.java

@@ -1387,6 +1387,8 @@ private final class CompleteSegmentCacheEntry implements SegmentCacheEntry
     private final SegmentCacheEntryIdentifier id;
     private final DataSegment dataSegment;
     private final String relativePathString;
+    @Nullable
+    private final SegmentId upgradedFromSegmentId;
     private SegmentLazyLoadFailCallback lazyLoadCallback = SegmentLazyLoadFailCallback.NOOP;
     private StorageLocation location;
     private File storageDir;
@@ -1401,6 +1403,7 @@ private CompleteSegmentCacheEntry(final DataSegment dataSegment)
       this.dataSegment = dataSegment;
       this.id = new SegmentCacheEntryIdentifier(dataSegment.getId());
       this.relativePathString = dataSegment.getId().toString();
+      this.upgradedFromSegmentId = dataSegment.getUpgradedFromSegmentId();
     }
 
     @Override
@@ -1479,7 +1482,12 @@ public void mount(StorageLocation mountLocation) throws SegmentLoadingException
             }
           }
           if (needsLoad) {
-            loadInLocationWithStartMarker(dataSegment, storageDir);
+            final File baseDir = findCachedBaseSegmentDir(mountLocation);
+            if (baseDir != null) {
+              copyFromBaseSegmentDir(baseDir, storageDir);
+            } else {
+              loadInLocationWithStartMarker(dataSegment, storageDir);
+            }
           }
           final SegmentizerFactory factory = getSegmentFactory(storageDir);
 
@@ -1663,6 +1671,71 @@ public File toPotentialLocation(final File location)
       return new File(location, relativePathString);
     }
 
+    /**
+     * Returns the cache directory of the base segment (the segment this one was upgraded from) on the given storage
+     * location if it is present and not corrupted, or null otherwise.
+     */
+    @GuardedBy("entryLock")
+    @Nullable
+    private File findCachedBaseSegmentDir(final StorageLocation mountLocation)
+    {
+      if (upgradedFromSegmentId == null) {
+        return null;
+      }
+      final File baseDir = new File(mountLocation.getPath(), upgradedFromSegmentId.toString());
+      if (baseDir.exists() && !isPossiblyCorrupted(baseDir)) {
+        return baseDir;
+      }
+      return null;
+    }
+
+    /**
+     * Copies files from the base segment's cache directory into the upgraded segment's cache directory,
+     * avoiding a deep-storage download. Uses a start-marker for crash safety, consistent with
+     * {@link #loadInLocationWithStartMarker}.
+     */
+    @GuardedBy("entryLock")
+    private void copyFromBaseSegmentDir(final File baseDir, final File destDir)
+        throws SegmentLoadingException
+    {
+      final File downloadStartMarker = new File(destDir, DOWNLOAD_START_MARKER_FILE_NAME);
+      try {
+        FileUtils.mkdirp(destDir);
+        if (!downloadStartMarker.createNewFile()) {
+          throw new SegmentLoadingException(
+              "Was not able to create new download marker for [%s]", destDir
+          );
+        }
+        log.info(
+            "Segment[%s] was upgraded from segment[%s] which is already cached; "
+            + "copying files from [%s] instead of downloading from deep storage.",
+            dataSegment.getId(),
+            upgradedFromSegmentId,
+            baseDir
+        );
+        final File[] files = baseDir.listFiles();
+        if (files != null) {
+          for (File file : files) {
+            if (!file.getName().equals(DOWNLOAD_START_MARKER_FILE_NAME)) {
+              Files.copy(file.toPath(), new File(destDir, file.getName()).toPath());
+            }
+          }
+        }
+        if (!downloadStartMarker.delete()) {
+          throw new SegmentLoadingException("Unable to remove marker file for [%s]", destDir);
+        }
+      }
+      catch (IOException e) {
+        throw new SegmentLoadingException(
+            e,
+            "Failed to copy base segment[%s] files from [%s] to [%s]",
+            upgradedFromSegmentId,
+            baseDir,
+            destDir
+        );
+      }
+    }
+
     @GuardedBy("entryLock")
     private void loadInLocationWithStartMarker(final DataSegment segment, final File storageDir)
         throws SegmentLoadingException