Source:
public long addSegment(String supervisorTaskId, String subTaskId, DataSegment segment, File segmentDir) throws IOException { final Iterator<StorageLocation> iterator = locationIterators.computeIfAbsent( supervisorTaskId, k -> { final Iterator<StorageLocation> cyclicIterator = Iterators.cycle(shuffleDataLocations); final int random = ThreadLocalRandom.current().nextInt(shuffleDataLocations.size()); IntStream.range(0, random).forEach(i -> cyclicIterator.next()); return cyclicIterator; } ); final File taskTempDir = taskConfig.getTaskTempDir(subTaskId); if (taskTempDir.mkdirs()) { taskTempDir.deleteOnExit(); } final File tempZippedFile = new File(taskTempDir, segment.getId().toString()); final long unzippedSizeBytes = CompressionUtils.zip(segmentDir, tempZippedFile, true); if (unzippedSizeBytes == 0) { throw new IOE( "Read 0 bytes from segmentDir[%s]", segmentDir.getAbsolutePath() ); } for (int i = 0; i < shuffleDataLocations.size(); i++) { final StorageLocation location = iterator.next(); final String partitionFilePath = getPartitionFilePath( supervisorTaskId, subTaskId, segment.getInterval(), segment.getShardSpec().getPartitionNum() ); final File destFile = location.reserve(partitionFilePath, segment.getId().toString(), tempZippedFile.length()); if (destFile != null) { try { FileUtils.forceMkdirParent(destFile); StreamUtils.retryCopy( Files.asByteSource(tempZippedFile), Files.asByteSink(destFile), t -> !(t instanceof InterruptedException) && !(t instanceof CancellationException) && (t instanceof Exception), DEFAULT_FILE_COPY_RETRY_COUNT ); if (!tempZippedFile.delete()) { log.warn("Couldn't delete file[%s]", tempZippedFile.getAbsolutePath()); } return unzippedSizeBytes; } catch (IOException e) { log.warn(e, "Failed to write segmentFile at [%s]", destFile); location.removeFile(tempZippedFile); } } } throw new ISE("Can't find location to handle segment[%s]", segment); }

Target:
long addSegment(String supervisorTaskId, String subTaskId, DataSegment segment, File segmentDir) throws IOException { final Iterator<StorageLocation> iterator = locationIterators.computeIfAbsent( supervisorTaskId, k -> { final Iterator<StorageLocation> cyclicIterator = Iterators.cycle(shuffleDataLocations); final int random = ThreadLocalRandom.current().nextInt(shuffleDataLocations.size()); IntStream.range(0, random).forEach(i -> cyclicIterator.next()); return cyclicIterator; } ); final File taskTempDir = taskConfig.getTaskTempDir(subTaskId); try (final Closer resourceCloser = Closer.create()) { if (taskTempDir.mkdirs()) { resourceCloser.register(() -> { try { FileUtils.forceDelete(taskTempDir); } catch (IOException e) { LOG.warn(e, "Failed to delete directory[%s]", taskTempDir.getAbsolutePath()); } }); } final File tempZippedFile = new File(taskTempDir, segment.getId().toString()); final long unzippedSizeBytes = CompressionUtils.zip(segmentDir, tempZippedFile); if (unzippedSizeBytes == 0) { throw new IOE( "Read 0 bytes from segmentDir[%s]", segmentDir.getAbsolutePath() ); } for (int i = 0; i < shuffleDataLocations.size(); i++) { final StorageLocation location = iterator.next(); final String partitionFilePath = getPartitionFilePath( supervisorTaskId, subTaskId, segment.getInterval(), segment.getShardSpec().getPartitionNum() ); final File destFile = location.reserve(partitionFilePath, segment.getId().toString(), tempZippedFile.length()); if (destFile != null) { FileUtils.forceMkdirParent(destFile); org.apache.druid.java.util.common.FileUtils.writeAtomically( destFile, out -> Files.asByteSource(tempZippedFile).copyTo(out) ); LOG.info( "Wrote intermediary segment for segment[%s] of subtask[%s] at [%s]", segment.getId(), subTaskId, destFile ); return unzippedSizeBytes; } } throw new ISE("Can't find location to handle segment[%s]", segment); } }