From 0644b73fabd3aef792c67d6436817e14e19a9f63 Mon Sep 17 00:00:00 2001 From: lentitude2tk Date: Tue, 23 Jun 2026 20:47:36 +0800 Subject: [PATCH] Improve volume upload interaction Signed-off-by: lentitude2tk --- .../milvus/bulkwriter/VolumeBulkWriter.java | 15 +- .../bulkwriter/VolumeBulkWriterParam.java | 18 + .../milvus/bulkwriter/VolumeFileManager.java | 593 +++++++++++++++--- .../bulkwriter/model/UploadProgress.java | 76 +++ .../request/volume/UploadFilesRequest.java | 117 ++++ .../bulkwriter/storage/StorageClient.java | 18 + .../storage/client/MinioStorageClient.java | 111 +++- .../milvus/bulkwriter/VolumeManagerTest.java | 342 ++++++++++ .../client/ProgressInputStreamTest.java | 68 ++ 9 files changed, 1249 insertions(+), 109 deletions(-) create mode 100644 sdk-bulkwriter/src/main/java/io/milvus/bulkwriter/model/UploadProgress.java create mode 100644 sdk-bulkwriter/src/test/java/io/milvus/bulkwriter/storage/client/ProgressInputStreamTest.java diff --git a/sdk-bulkwriter/src/main/java/io/milvus/bulkwriter/VolumeBulkWriter.java b/sdk-bulkwriter/src/main/java/io/milvus/bulkwriter/VolumeBulkWriter.java index af716711f..2cda4569f 100644 --- a/sdk-bulkwriter/src/main/java/io/milvus/bulkwriter/VolumeBulkWriter.java +++ b/sdk-bulkwriter/src/main/java/io/milvus/bulkwriter/VolumeBulkWriter.java @@ -21,7 +21,6 @@ import com.google.common.collect.Lists; import com.google.gson.JsonObject; -import io.milvus.bulkwriter.common.clientenum.ConnectType; import io.milvus.bulkwriter.model.UploadFilesResult; import io.milvus.bulkwriter.request.volume.UploadFilesRequest; import io.milvus.common.utils.ExceptionUtils; @@ -64,7 +63,7 @@ public VolumeBulkWriter(VolumeBulkWriterParam bulkWriterParam) throws IOExceptio private VolumeFileManager initVolumeFileManagerParams(VolumeBulkWriterParam bulkWriterParam) throws IOException { VolumeFileManagerParam volumeFileManagerParam = VolumeFileManagerParam.newBuilder() .withCloudEndpoint(bulkWriterParam.getCloudEndpoint()).withApiKey(bulkWriterParam.getApiKey()) - .withVolumeName(bulkWriterParam.getVolumeName()).withConnectType(ConnectType.AUTO) + .withVolumeName(bulkWriterParam.getVolumeName()).withConnectType(bulkWriterParam.getConnectType()) .build(); return new VolumeFileManager(volumeFileManagerParam); } @@ -149,9 +148,13 @@ protected void callBack(List fileList) { @Override public void close() throws Exception { - logger.info("execute remaining actions to prevent loss of memory data or residual empty directories."); - exit(); - logger.info(String.format("RemoteBulkWriter done! output remote files: %s", getBatchFiles())); + try { + logger.info("execute remaining actions to prevent loss of memory data or residual empty directories."); + exit(); + logger.info(String.format("RemoteBulkWriter done! output remote files: %s", getBatchFiles())); + } finally { + volumeFileManager.shutdownGracefully(); + } } private void serialImportData(List fileList) { @@ -199,6 +202,6 @@ private static String getMinioFilePath(String remotePath, String relativeFilePat relativeFilePath = relativeFilePath.startsWith("/") ? relativeFilePath.substring(1) : relativeFilePath; Path relative = Paths.get(relativeFilePath); Path joinedPath = remote.resolve(relative); - return joinedPath.toString(); + return joinedPath.toString().replace("\\", "/"); } } diff --git a/sdk-bulkwriter/src/main/java/io/milvus/bulkwriter/VolumeBulkWriterParam.java b/sdk-bulkwriter/src/main/java/io/milvus/bulkwriter/VolumeBulkWriterParam.java index 71521ed23..7647e4753 100644 --- a/sdk-bulkwriter/src/main/java/io/milvus/bulkwriter/VolumeBulkWriterParam.java +++ b/sdk-bulkwriter/src/main/java/io/milvus/bulkwriter/VolumeBulkWriterParam.java @@ -20,6 +20,7 @@ package io.milvus.bulkwriter; import io.milvus.bulkwriter.common.clientenum.BulkFileType; +import io.milvus.bulkwriter.common.clientenum.ConnectType; import io.milvus.bulkwriter.common.utils.V2AdapterUtils; import io.milvus.exception.ParamException; import io.milvus.param.ParamUtils; @@ -42,6 +43,7 @@ public class VolumeBulkWriterParam { private final String cloudEndpoint; private final String apiKey; private final String volumeName; + private final ConnectType connectType; private VolumeBulkWriterParam(Builder builder) { this.collectionSchema = builder.collectionSchema; @@ -53,6 +55,7 @@ private VolumeBulkWriterParam(Builder builder) { this.cloudEndpoint = builder.cloudEndpoint; this.apiKey = builder.apiKey; this.volumeName = builder.volumeName; + this.connectType = builder.connectType; } public CreateCollectionReq.CollectionSchema getCollectionSchema() { @@ -87,6 +90,10 @@ public String getVolumeName() { return volumeName; } + public ConnectType getConnectType() { + return connectType; + } + @Override public String toString() { return "VolumeBulkWriterParam{" + @@ -96,6 +103,7 @@ public String toString() { ", fileType=" + fileType + ", cloudEndpoint='" + cloudEndpoint + '\'' + ", volumeName='" + volumeName + '\'' + + ", connectType=" + connectType + '}'; } @@ -118,6 +126,8 @@ public static final class Builder { private String volumeName; + private ConnectType connectType = ConnectType.AUTO; + private Builder() { } @@ -184,6 +194,11 @@ public Builder withVolumeName(String volumeName) { return this; } + public Builder withConnectType(ConnectType connectType) { + this.connectType = connectType; + return this; + } + /** * Verifies parameters and creates a new {@link VolumeBulkWriterParam} instance. * @@ -195,6 +210,9 @@ public VolumeBulkWriterParam build() throws ParamException { if (collectionSchema == null) { throw new ParamException("collectionSchema cannot be null"); } + ParamUtils.CheckNullEmptyString(cloudEndpoint, "cloudEndpoint"); + ParamUtils.CheckNullEmptyString(apiKey, "apiKey"); + ParamUtils.CheckNullEmptyString(volumeName, "volumeName"); return new VolumeBulkWriterParam(this); } diff --git a/sdk-bulkwriter/src/main/java/io/milvus/bulkwriter/VolumeFileManager.java b/sdk-bulkwriter/src/main/java/io/milvus/bulkwriter/VolumeFileManager.java index 02a2a4b7e..cbf20eb7c 100644 --- a/sdk-bulkwriter/src/main/java/io/milvus/bulkwriter/VolumeFileManager.java +++ b/sdk-bulkwriter/src/main/java/io/milvus/bulkwriter/VolumeFileManager.java @@ -23,6 +23,7 @@ import io.milvus.bulkwriter.common.clientenum.ConnectType; import io.milvus.bulkwriter.common.utils.FileUtils; import io.milvus.bulkwriter.model.UploadFilesResult; +import io.milvus.bulkwriter.model.UploadProgress; import io.milvus.bulkwriter.request.volume.ApplyVolumeRequest; import io.milvus.bulkwriter.request.volume.UploadFilesRequest; import io.milvus.bulkwriter.resolver.EndpointResolver; @@ -31,38 +32,41 @@ import io.milvus.bulkwriter.storage.StorageClient; import io.milvus.bulkwriter.storage.client.MinioStorageClient; import io.milvus.exception.ParamException; +import io.minio.errors.ErrorResponseException; import org.apache.commons.lang3.StringUtils; import org.apache.commons.lang3.tuple.Pair; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import java.io.File; +import java.io.IOException; import java.nio.file.Path; import java.nio.file.Paths; import java.time.Instant; -import java.util.Date; +import java.util.HashMap; +import java.util.HashSet; import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Set; import java.util.concurrent.*; -import java.util.concurrent.atomic.AtomicInteger; -import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; public class VolumeFileManager { private static final Logger logger = LoggerFactory.getLogger(VolumeFileManager.class); + private static final long DEFAULT_CREDENTIAL_REFRESH_MARGIN_SECONDS = 300L; private final String cloudEndpoint; private final String apiKey; private final String volumeName; private final ConnectType connectType; - private final ExecutorService executor; - private StorageClient storageClient; - private ApplyVolumeResponse applyVolumeResponse; + private volatile VolumeSession lastVolumeSession; public VolumeFileManager(VolumeFileManagerParam volumeFileManagerParam) { this.cloudEndpoint = volumeFileManagerParam.getCloudEndpoint(); this.apiKey = volumeFileManagerParam.getApiKey(); this.volumeName = volumeFileManagerParam.getVolumeName(); - this.connectType = volumeFileManagerParam.getConnectType(); - this.executor = Executors.newFixedThreadPool(10); + this.connectType = volumeFileManagerParam.getConnectType() == null ? ConnectType.AUTO : volumeFileManagerParam.getConnectType(); } /** @@ -78,77 +82,137 @@ public CompletableFuture uploadFilesAsync(UploadFilesRequest String localDirOrFilePath = request.getSourceFilePath(); Pair, Long> localPathPair = FileUtils.processLocalPath(localDirOrFilePath); String volumePath = convertDirPath(request.getTargetVolumePath()); - - refreshVolumeAndClient(volumePath); - initValidator(localPathPair); - - AtomicInteger currentFileCount = new AtomicInteger(0); - AtomicLong processedBytes = new AtomicLong(0); + int uploadConcurrency = Math.max(1, request.getUploadConcurrency()); + int maxRetries = Math.max(0, request.getMaxRetries()); + long retryIntervalMillis = Math.max(0L, request.getRetryIntervalMillis()); + long partSizeBytes = Math.max(0L, request.getPartSizeBytes()); long totalBytes = localPathPair.getValue(); long totalFilesCount = localPathPair.getKey().size(); long startTime = System.currentTimeMillis(); + logger.info("Starting volume upload: sourcePath:{}, volumeName:{}, volumePath:{}, totalFileCount:{}, totalFileSize:{} bytes ({}), uploadConcurrency:{}, maxRetries:{}, retryInterval:{}, partSize:{}, startTime:{}", + localDirOrFilePath, volumeName, volumePath, totalFilesCount, totalBytes, formatBytes(totalBytes), + uploadConcurrency, maxRetries, formatDurationMillis(retryIntervalMillis), + formatPartSize(partSizeBytes), Instant.ofEpochMilli(startTime)); + + VolumeSession initialSession; + try { + initialSession = refreshVolumeAndClient(volumePath); + } catch (RuntimeException e) { + logUploadFailed(localDirOrFilePath, volumePath, startTime); + throw e; + } + try { + initValidator(localPathPair, initialSession.applyVolumeResponse); + } catch (RuntimeException e) { + closeVolumeSession(initialSession); + logUploadFailed(localDirOrFilePath, volumePath, startTime); + throw e; + } + UploadContext uploadContext = new UploadContext(initialSession); + ExecutorService uploadExecutor = Executors.newFixedThreadPool(uploadConcurrency); + UploadProgressTracker progressTracker = new UploadProgressTracker(totalBytes, totalFilesCount, request.getProgressListener()); return CompletableFuture.allOf(localPathPair.getKey().stream() .map(localFilePath -> CompletableFuture.runAsync(() -> { File file = new File(localFilePath); + String progressFilePath = file.getAbsolutePath(); long fileStartTime = System.currentTimeMillis(); try { - uploadLocalFileToVolume(localFilePath, localDirOrFilePath, volumePath); - long bytes = processedBytes.addAndGet(file.length()); - int completeCount = currentFileCount.incrementAndGet(); + uploadLocalFileToVolume(localFilePath, localDirOrFilePath, volumePath, maxRetries, retryIntervalMillis, progressTracker, uploadContext, partSizeBytes); + UploadProgressSnapshot progress = progressTracker.finishFile(progressFilePath, file.length()); long elapsed = System.currentTimeMillis() - fileStartTime; - double percent = totalBytes == 0 ? 100.0 : (bytes * 100.0 / totalBytes); - logger.info("Uploaded file {}/{}: {} ({} bytes) elapsed:{} ms, progress(total bytes): {}/{} bytes, progress(total percentage):{}%", - completeCount, totalFilesCount, localFilePath, file.length(), elapsed, bytes, totalBytes, String.format("%.2f", percent)); + logger.info("Uploaded file {}/{}: {} ({} bytes) elapsed:{} ms, progress(total bytes): {}/{} bytes, progress(total percentage):{}%, speedBPS:{}, estimatedRemainingTime:{}", + progress.completedFiles, totalFilesCount, localFilePath, file.length(), elapsed, + progress.uploadedBytes, totalBytes, String.format("%.2f", progress.percent), + progress.speedBps, progress.estimatedRemainingTime); } catch (Exception e) { logger.error("Upload failed: {}", localFilePath, e); throw new CompletionException(e); } - }, executor)).toArray(CompletableFuture[]::new)) + }, uploadExecutor)).toArray(CompletableFuture[]::new)) .whenComplete((v, t) -> { + uploadExecutor.shutdown(); + uploadContext.closeSessions(); + if (t != null) { + logUploadFailed(localDirOrFilePath, volumePath, startTime); + } }) .thenApply(v -> { - long totalElapsed = (System.currentTimeMillis() - startTime) / 1000; - logger.info("all files in {} has been async uploaded to volume, volumeName:{}, volumePath:{}, totalFileCount:{}, totalFileSize:{}, cost times:{} s", - localDirOrFilePath, applyVolumeResponse.getVolumeName(), volumePath, localPathPair.getKey().size(), localPathPair.getValue(), totalElapsed); + progressTracker.finishUpload(); + long endTime = System.currentTimeMillis(); + long totalElapsed = endTime - startTime; + VolumeSession session = uploadContext.currentSession(); + logger.info("Volume upload completed: sourcePath:{}, volumeName:{}, volumePath:{}, totalFileCount:{}, totalFileSize:{} bytes ({}), endTime:{}, totalElapsed:{}", + localDirOrFilePath, session.applyVolumeResponse.getVolumeName(), volumePath, + localPathPair.getKey().size(), localPathPair.getValue(), formatBytes(localPathPair.getValue()), + Instant.ofEpochMilli(endTime), formatDurationMillis(totalElapsed)); return UploadFilesResult.builder() - .volumeName(applyVolumeResponse.getVolumeName()) + .volumeName(session.applyVolumeResponse.getVolumeName()) .path(volumePath) .build(); }); } /** - * Gracefully shuts down the internal thread pool executor. - *

- * This method attempts to stop accepting new tasks and waits for existing - * tasks to complete within a timeout period. If tasks do not finish within - * the timeout, it will forcefully shut down the executor. - *

- *

- * Usage recommendation: - *

    - *
  • Call this method when the VolumeFileManager is no longer needed.
    • - *
- *

- * Thread interruption is respected, and the interrupt status is restored if interrupted during shutdown. + * Synchronously uploads a local file or directory to the specified path within the Volume. */ + public UploadFilesResult uploadFiles(UploadFilesRequest request) throws ExecutionException, InterruptedException { + return uploadFilesAsync(request).get(); + } + public void shutdownGracefully() { - executor.shutdown(); - try { - if (!executor.awaitTermination(2, TimeUnit.SECONDS)) { - logger.warn("Executor didn't terminate in time, forcing shutdown..."); - executor.shutdownNow(); - } - } catch (InterruptedException e) { - logger.error("Interrupted while waiting for executor to shutdown", e); - executor.shutdownNow(); - Thread.currentThread().interrupt(); + VolumeSession session = lastVolumeSession; + lastVolumeSession = null; + closeVolumeSession(session); + } + + private void logUploadFailed(String localDirOrFilePath, String volumePath, long startTime) { + long endTime = System.currentTimeMillis(); + logger.warn("Volume upload failed: sourcePath:{}, volumeName:{}, volumePath:{}, endTime:{}, totalElapsed:{}", + localDirOrFilePath, volumeName, volumePath, + Instant.ofEpochMilli(endTime), formatDurationMillis(endTime - startTime)); + } + + private static String formatPartSize(long partSizeBytes) { + if (partSizeBytes <= 0L) { + return "auto"; + } + return partSizeBytes + " bytes (" + formatBytes(partSizeBytes) + ")"; + } + + private static String formatBytes(long bytes) { + if (bytes < 1024L) { + return bytes + " B"; } + double value = bytes; + String[] units = new String[]{"B", "KiB", "MiB", "GiB", "TiB", "PiB"}; + int unitIndex = 0; + while (value >= 1024.0 && unitIndex < units.length - 1) { + value /= 1024.0; + unitIndex++; + } + return String.format(Locale.ROOT, "%.2f %s", value, units[unitIndex]); } - private void initValidator(Pair, Long> localPathPair) { + private static String formatDurationMillis(long durationMillis) { + if (durationMillis < 0L) { + return "unknown"; + } + long totalSeconds = (durationMillis + 999L) / 1000L; + long hours = totalSeconds / 3600L; + long minutes = (totalSeconds % 3600L) / 60L; + long seconds = totalSeconds % 60L; + if (hours > 0L) { + return String.format(Locale.ROOT, "%dh %02dm %02ds", hours, minutes, seconds); + } + if (minutes > 0L) { + return String.format(Locale.ROOT, "%dm %02ds", minutes, seconds); + } + return String.format(Locale.ROOT, "%ds", seconds); + } + + private void initValidator(Pair, Long> localPathPair, ApplyVolumeResponse applyVolumeResponse) { Long maxContentLength = applyVolumeResponse.getCondition().getMaxContentLength(); Long uploadFileContentLength = localPathPair.getValue(); if (uploadFileContentLength > maxContentLength) { @@ -172,40 +236,69 @@ private void initValidator(Pair, Long> localPathPair) { } } - private void refreshVolumeAndClient(String path) { + private VolumeSession refreshVolumeAndClient(String path) { logger.info("refreshing Volume info..."); ApplyVolumeRequest applyVolumeRequest = ApplyVolumeRequest.builder() .apiKey(apiKey) .volumeName(volumeName) .path(path) .build(); - String result = DataVolumeUtils.applyVolume(cloudEndpoint, applyVolumeRequest); - applyVolumeResponse = new Gson().fromJson(result, ApplyVolumeResponse.class); + String result = applyVolume(applyVolumeRequest); + ApplyVolumeResponse applyVolumeResponse = new Gson().fromJson(result, ApplyVolumeResponse.class); logger.info("volume info refreshed"); + StorageClient storageClient = createStorageClient(applyVolumeResponse); + VolumeSession session = new VolumeSession(applyVolumeResponse, storageClient, + Instant.parse(applyVolumeResponse.getCredentials().getExpireTime())); + lastVolumeSession = session; + logger.info("storage client refreshed"); + return session; + } + + protected String applyVolume(ApplyVolumeRequest applyVolumeRequest) { + return DataVolumeUtils.applyVolume(cloudEndpoint, applyVolumeRequest); + } + + protected StorageClient createStorageClient(ApplyVolumeResponse applyVolumeResponse) { String endpoint = EndpointResolver.resolveEndpoint(applyVolumeResponse.getEndpoint(), applyVolumeResponse.getCloud(), applyVolumeResponse.getRegion(), connectType); - storageClient = MinioStorageClient.getStorageClient( + return MinioStorageClient.getStorageClient( applyVolumeResponse.getCloud(), endpoint, applyVolumeResponse.getCredentials().getTmpAK(), applyVolumeResponse.getCredentials().getTmpSK(), applyVolumeResponse.getCredentials().getSessionToken(), applyVolumeResponse.getRegion(), null); - logger.info("storage client refreshed"); } private String convertDirPath(String inputPath) { if (StringUtils.isEmpty(inputPath) || inputPath.equals("/")) { return ""; } - if (inputPath.endsWith("/")) { - return inputPath; + String[] parts = inputPath.replace("\\", "/").split("/"); + StringBuilder builder = new StringBuilder(); + for (String part : parts) { + if (StringUtils.isEmpty(part) || part.equals(".")) { + continue; + } + if (part.equals("..")) { + throw new ParamException("target volume path must not escape the volume root: " + inputPath); + } + if (builder.length() > 0) { + builder.append("/"); + } + builder.append(part); + } + if (builder.length() == 0) { + return ""; } - return inputPath + "/"; + return builder + "/"; } - private void uploadLocalFileToVolume(String localFilePath, String rootPath, String volumePath) { + private void uploadLocalFileToVolume(String localFilePath, String rootPath, String volumePath, + int maxRetries, long retryIntervalMillis, + UploadProgressTracker progressTracker, + UploadContext uploadContext, long partSizeBytes) { File file = new File(localFilePath); Path filePath = file.toPath().toAbsolutePath(); Path root = Paths.get(rootPath).toAbsolutePath(); @@ -217,65 +310,379 @@ private void uploadLocalFileToVolume(String localFilePath, String rootPath, Stri relativePath = root.relativize(filePath).toString().replace("\\", "/"); } - String remoteFilePath = applyVolumeResponse.getVolumePrefix() + volumePath + relativePath; - putObjectWithRetry(file, remoteFilePath, volumePath); + VolumeSession session = uploadContext.currentSession(); + String remoteFilePath = session.applyVolumeResponse.getVolumePrefix() + volumePath + relativePath; + putObjectWithRetry(file, remoteFilePath, volumePath, maxRetries, retryIntervalMillis, progressTracker, uploadContext, partSizeBytes); } - private void putObjectWithRetry(File file, String remoteFilePath, String volumePath) { - refreshIfExpire(volumePath); + private void putObjectWithRetry(File file, String remoteFilePath, String volumePath, + int maxRetries, long retryIntervalMillis, + UploadProgressTracker progressTracker, + UploadContext uploadContext, long partSizeBytes) { String msg = "upload " + file.getAbsolutePath(); + FileUploadProgress progress = new FileUploadProgress(progressTracker, file.getAbsolutePath(), file.length()); withRetry(msg, () -> { - try { - storageClient.putObject(file, applyVolumeResponse.getBucketName(), remoteFilePath); - } catch (Exception e) { - throw new RuntimeException(e); - } - }, volumePath); + progress.reset(); + VolumeSession session = refreshIfExpire(volumePath, uploadContext); + session.storageClient.putObject(file, session.applyVolumeResponse.getBucketName(), remoteFilePath, progress, partSizeBytes); + return null; + }, volumePath, maxRetries, retryIntervalMillis, uploadContext); } - private void refreshIfExpire(String volumePath) { - Instant instant = Instant.parse(applyVolumeResponse.getCredentials().getExpireTime()); - Date expireTime = Date.from(instant); - if (new Date().after(expireTime)) { - synchronized (this) { - if (new Date().after(expireTime)) { - refreshVolumeAndClient(volumePath); - } + private VolumeSession refreshIfExpire(String volumePath, UploadContext uploadContext) { + VolumeSession session = uploadContext.currentSession(); + Instant refreshAt = session.expireTime.minusSeconds(DEFAULT_CREDENTIAL_REFRESH_MARGIN_SECONDS); + if (Instant.now().isBefore(refreshAt)) { + return session; + } + synchronized (uploadContext.refreshLock) { + session = uploadContext.currentSession(); + refreshAt = session.expireTime.minusSeconds(DEFAULT_CREDENTIAL_REFRESH_MARGIN_SECONDS); + if (Instant.now().isBefore(refreshAt)) { + return session; } + return refreshUploadContext(volumePath, uploadContext); + } + } + + private VolumeSession refreshUploadContext(String volumePath, UploadContext uploadContext) { + VolumeSession session = refreshVolumeAndClient(volumePath); + uploadContext.setSession(session); + return session; + } + + private static void closeVolumeSession(VolumeSession session) { + if (session == null) { + return; } + session.storageClient.close(); } - private T withRetry(String actionName, Callable callable, String volumePath) { - final int maxRetries = 5; - int attempt = 0; - while (attempt < maxRetries) { + private T withRetry(String actionName, Callable callable, String volumePath, + int maxRetries, long retryIntervalMillis, + UploadContext uploadContext) { + int failedAttempts = 0; + while (true) { try { return callable.call(); - } catch (RuntimeException e) { - throw e; } catch (Exception e) { - attempt++; - refreshVolumeAndClient(volumePath); - logger.warn("Attempt {} failed to {}", attempt, actionName, e); - if (attempt == maxRetries) { - throw new RuntimeException(actionName + " failed after " + maxRetries + " attempts", e); + if (hasCause(e, ProgressCallbackException.class)) { + throw new RuntimeException("Upload progress callback failed", e); + } + if (!isRetryableException(e)) { + if (e instanceof RuntimeException) { + throw (RuntimeException) e; + } + throw new RuntimeException(actionName + " failed", e); } + failedAttempts++; + logger.warn("Attempt {} failed to {}", failedAttempts, actionName, e); + if (failedAttempts > maxRetries) { + throw new RuntimeException(actionName + " failed after " + failedAttempts + " attempts", e); + } + refreshUploadContext(volumePath, uploadContext); try { - Thread.sleep(5000L); - } catch (InterruptedException ignored) { + Thread.sleep(retryIntervalMillis); + } catch (InterruptedException interruptedException) { + Thread.currentThread().interrupt(); + throw new RuntimeException(actionName + " interrupted while waiting to retry", interruptedException); } } } - throw new RuntimeException(actionName + " failed unexpectedly."); } - private void withRetry(String actionName, Runnable runnable, String volumePath) { + private void withRetry(String actionName, Runnable runnable, String volumePath, + int maxRetries, long retryIntervalMillis, + UploadContext uploadContext) { withRetry(actionName, () -> { runnable.run(); return null; - }, volumePath); + }, volumePath, maxRetries, retryIntervalMillis, uploadContext); } + private boolean hasCause(Throwable throwable, Class causeClass) { + Throwable current = throwable; + while (current != null) { + if (causeClass.isInstance(current)) { + return true; + } + current = current.getCause(); + } + return false; + } + + private boolean isRetryableException(Throwable throwable) { + if (hasCause(throwable, ParamException.class) || hasCause(throwable, IllegalArgumentException.class)) { + return false; + } + ErrorResponseException errorResponseException = findCause(throwable, ErrorResponseException.class); + if (errorResponseException != null) { + return isRetryableS3Error(errorResponseException); + } + return hasCause(throwable, IOException.class) || hasCause(throwable, TimeoutException.class); + } + + private boolean isRetryableS3Error(ErrorResponseException exception) { + int statusCode = exception.response() == null ? 0 : exception.response().code(); + if (statusCode == 408 || statusCode == 429 || statusCode >= 500) { + return true; + } + String code = exception.errorResponse() == null ? "" : exception.errorResponse().code(); + return "RequestTimeout".equals(code) + || "SlowDown".equals(code) + || "InternalError".equals(code) + || "ServiceUnavailable".equals(code) + || "Throttling".equals(code) + || "ThrottlingException".equals(code) + || "RequestLimitExceeded".equals(code); + } + + private T findCause(Throwable throwable, Class causeClass) { + Throwable current = throwable; + while (current != null) { + if (causeClass.isInstance(current)) { + return causeClass.cast(current); + } + current = current.getCause(); + } + return null; + } + + private VolumeSession currentSession() { + VolumeSession session = lastVolumeSession; + if (session == null) { + throw new IllegalStateException("Volume session is not initialized"); + } + return session; + } + + private static class UploadContext { + private final AtomicReference sessionRef; + private final Object refreshLock = new Object(); + private final Set sessions = ConcurrentHashMap.newKeySet(); + + private UploadContext(VolumeSession session) { + this.sessionRef = new AtomicReference<>(session); + this.sessions.add(session); + } + + private VolumeSession currentSession() { + VolumeSession session = sessionRef.get(); + if (session == null) { + throw new IllegalStateException("Volume session is not initialized"); + } + return session; + } + + private void setSession(VolumeSession session) { + sessionRef.set(session); + sessions.add(session); + } + + private void closeSessions() { + for (VolumeSession session : sessions) { + closeVolumeSession(session); + } + sessions.clear(); + } + } + + private static class UploadProgressTracker { + private static final long LOG_INTERVAL_MILLIS = 5000L; + private static final double LOG_PERCENT_STEP = 1.0; + + private final long totalBytes; + private final long totalFiles; + private final UploadFilesRequest.ProgressListener progressListener; + private final long startTimeMillis; + private final Map fileProgress = new HashMap<>(); + private final Set completedFiles = new HashSet<>(); + private long uploadedBytes = 0L; + private long lastLogTimeMillis = 0L; + private double lastLoggedPercent = -1.0; + + private UploadProgressTracker(long totalBytes, long totalFiles, + UploadFilesRequest.ProgressListener progressListener) { + this.totalBytes = totalBytes; + this.totalFiles = totalFiles; + this.progressListener = progressListener; + this.startTimeMillis = System.currentTimeMillis(); + } + + synchronized void resetFile(String filePath) { + long previous = fileProgress.getOrDefault(filePath, 0L); + uploadedBytes -= previous; + fileProgress.put(filePath, 0L); + } + + void updateFile(String filePath, long fileSize, long chunkBytes) { + UploadProgress progress = null; + synchronized (this) { + if (chunkBytes <= 0) { + return; + } + long previous = fileProgress.getOrDefault(filePath, 0L); + long current = Math.min(fileSize, previous + chunkBytes); + long delta = current - previous; + if (delta <= 0) { + return; + } + fileProgress.put(filePath, current); + uploadedBytes += delta; + progress = progressIfNeeded(filePath, current, fileSize); + } + if (progress != null) { + emitProgress(progress); + } + } + + UploadProgressSnapshot finishFile(String filePath, long fileSize) { + UploadProgress progress; + UploadProgressSnapshot snapshot; + synchronized (this) { + long previous = fileProgress.getOrDefault(filePath, 0L); + long current = Math.max(previous, fileSize); + fileProgress.put(filePath, current); + uploadedBytes += current - previous; + completedFiles.add(filePath); + double percent = percent(); + progress = snapshot(filePath, current, fileSize, percent); + long speedBps = speedBps(uploadedBytes, System.currentTimeMillis()); + snapshot = new UploadProgressSnapshot(uploadedBytes, completedFiles.size(), percent, + speedBps, formatDurationMillis(estimatedRemainingTimeMillis(uploadedBytes, speedBps))); + markProgressEmitted(percent); + } + emitProgress(progress); + return snapshot; + } + + void finishUpload() { + UploadProgress progress; + synchronized (this) { + progress = snapshot("", 0L, 0L, percent()); + markProgressEmitted(progress.getPercent()); + } + emitProgress(progress); + } + + private double percent() { + if (totalBytes == 0) { + return 100.0; + } + return Math.min(100.0, uploadedBytes * 100.0 / totalBytes); + } + + private long speedBps(long currentUploadedBytes, long nowMillis) { + long elapsedMillis = Math.max(1L, nowMillis - startTimeMillis); + return (long) (currentUploadedBytes * 1000.0 / elapsedMillis); + } + + private long estimatedRemainingTimeMillis(long currentUploadedBytes, long speedBps) { + long remainingBytes = Math.max(0L, totalBytes - currentUploadedBytes); + if (remainingBytes == 0L) { + return 0L; + } + if (speedBps <= 0L) { + return -1L; + } + return (long) Math.ceil(remainingBytes * 1000.0 / speedBps); + } + + private UploadProgress progressIfNeeded(String currentFile, long currentFileUploadedBytes, long currentFileTotalBytes) { + long now = System.currentTimeMillis(); + double percent = percent(); + if (percent - lastLoggedPercent >= LOG_PERCENT_STEP || now - lastLogTimeMillis >= LOG_INTERVAL_MILLIS) { + lastLogTimeMillis = now; + lastLoggedPercent = percent; + return snapshot(currentFile, currentFileUploadedBytes, currentFileTotalBytes, percent); + } + return null; + } + + private UploadProgress snapshot(String currentFile, long currentFileUploadedBytes, + long currentFileTotalBytes, double percent) { + return new UploadProgress(uploadedBytes, totalBytes, completedFiles.size(), totalFiles, + currentFile, currentFileUploadedBytes, currentFileTotalBytes, percent); + } + + private void markProgressEmitted(double percent) { + lastLogTimeMillis = System.currentTimeMillis(); + lastLoggedPercent = percent; + } + + private void emitProgress(UploadProgress progress) { + long speedBps = speedBps(progress.getUploadedBytes(), System.currentTimeMillis()); + String estimatedRemainingTime = formatDurationMillis(estimatedRemainingTimeMillis(progress.getUploadedBytes(), speedBps)); + logger.info("Upload progress: {}/{} bytes, progress:{}%, files:{}/{}, speedBPS:{}, estimatedRemainingTime:{}", + progress.getUploadedBytes(), progress.getTotalBytes(), + String.format("%.2f", progress.getPercent()), + progress.getCompletedFiles(), progress.getTotalFiles(), + speedBps, estimatedRemainingTime); + if (progressListener != null) { + try { + progressListener.onProgress(progress); + } catch (RuntimeException e) { + throw new ProgressCallbackException(e); + } + } + } + } + + private static class ProgressCallbackException extends RuntimeException { + private ProgressCallbackException(Throwable cause) { + super(cause); + } + } + + private static class FileUploadProgress implements StorageClient.UploadProgressListener { + private final UploadProgressTracker progressTracker; + private final String filePath; + private final long fileSize; + + private FileUploadProgress(UploadProgressTracker progressTracker, String filePath, long fileSize) { + this.progressTracker = progressTracker; + this.filePath = filePath; + this.fileSize = fileSize; + } + + private void reset() { + progressTracker.resetFile(filePath); + } + + @Override + public void onProgress(long bytes) { + progressTracker.updateFile(filePath, fileSize, bytes); + } + } + + private static class UploadProgressSnapshot { + private final long uploadedBytes; + private final int completedFiles; + private final double percent; + private final long speedBps; + private final String estimatedRemainingTime; + + private UploadProgressSnapshot(long uploadedBytes, int completedFiles, double percent, + long speedBps, String estimatedRemainingTime) { + this.uploadedBytes = uploadedBytes; + this.completedFiles = completedFiles; + this.percent = percent; + this.speedBps = speedBps; + this.estimatedRemainingTime = estimatedRemainingTime; + } + } + + private static class VolumeSession { + private final ApplyVolumeResponse applyVolumeResponse; + private final StorageClient storageClient; + private final Instant expireTime; + + private VolumeSession(ApplyVolumeResponse applyVolumeResponse, StorageClient storageClient, Instant expireTime) { + this.applyVolumeResponse = applyVolumeResponse; + this.storageClient = storageClient; + this.expireTime = expireTime; + } + } -} \ No newline at end of file +} diff --git a/sdk-bulkwriter/src/main/java/io/milvus/bulkwriter/model/UploadProgress.java b/sdk-bulkwriter/src/main/java/io/milvus/bulkwriter/model/UploadProgress.java new file mode 100644 index 000000000..df9739d09 --- /dev/null +++ b/sdk-bulkwriter/src/main/java/io/milvus/bulkwriter/model/UploadProgress.java @@ -0,0 +1,76 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package io.milvus.bulkwriter.model; + +public class UploadProgress { + private final long uploadedBytes; + private final long totalBytes; + private final int completedFiles; + private final long totalFiles; + private final String currentFile; + private final long currentFileUploadedBytes; + private final long currentFileTotalBytes; + private final double percent; + + public UploadProgress(long uploadedBytes, long totalBytes, int completedFiles, long totalFiles, + String currentFile, long currentFileUploadedBytes, + long currentFileTotalBytes, double percent) { + this.uploadedBytes = uploadedBytes; + this.totalBytes = totalBytes; + this.completedFiles = completedFiles; + this.totalFiles = totalFiles; + this.currentFile = currentFile; + this.currentFileUploadedBytes = currentFileUploadedBytes; + this.currentFileTotalBytes = currentFileTotalBytes; + this.percent = percent; + } + + public long getUploadedBytes() { + return uploadedBytes; + } + + public long getTotalBytes() { + return totalBytes; + } + + public int getCompletedFiles() { + return completedFiles; + } + + public long getTotalFiles() { + return totalFiles; + } + + public String getCurrentFile() { + return currentFile; + } + + public long getCurrentFileUploadedBytes() { + return currentFileUploadedBytes; + } + + public long getCurrentFileTotalBytes() { + return currentFileTotalBytes; + } + + public double getPercent() { + return percent; + } +} diff --git a/sdk-bulkwriter/src/main/java/io/milvus/bulkwriter/request/volume/UploadFilesRequest.java b/sdk-bulkwriter/src/main/java/io/milvus/bulkwriter/request/volume/UploadFilesRequest.java index 5b8348227..ec7aa21a6 100644 --- a/sdk-bulkwriter/src/main/java/io/milvus/bulkwriter/request/volume/UploadFilesRequest.java +++ b/sdk-bulkwriter/src/main/java/io/milvus/bulkwriter/request/volume/UploadFilesRequest.java @@ -19,6 +19,8 @@ package io.milvus.bulkwriter.request.volume; +import io.milvus.bulkwriter.model.UploadProgress; + public class UploadFilesRequest { /** * The full path of a local file or directory: @@ -34,6 +36,31 @@ public class UploadFilesRequest { */ private String targetVolumePath; + /** + * The maximum number of files to upload concurrently. + */ + private int uploadConcurrency = 5; + + /** + * The maximum retry count for each file. + */ + private int maxRetries = 5; + + /** + * Retry interval in milliseconds. + */ + private long retryIntervalMillis = 5000L; + + /** + * Optional callback for upload progress snapshots. + */ + private ProgressListener progressListener = null; + + /** + * Multipart upload part size in bytes. Zero or negative means automatic. + */ + private long partSizeBytes = 0L; + public UploadFilesRequest() { } @@ -45,6 +72,11 @@ public UploadFilesRequest(String sourceFilePath, String targetVolumePath) { protected UploadFilesRequest(UploadFilesRequestBuilder builder) { this.sourceFilePath = builder.sourceFilePath; this.targetVolumePath = builder.targetVolumePath; + this.uploadConcurrency = builder.uploadConcurrency; + this.maxRetries = builder.maxRetries; + this.retryIntervalMillis = builder.retryIntervalMillis; + this.progressListener = builder.progressListener; + this.partSizeBytes = builder.partSizeBytes; } public String getSourceFilePath() { @@ -63,11 +95,56 @@ public void setTargetVolumePath(String targetVolumePath) { this.targetVolumePath = targetVolumePath; } + public int getUploadConcurrency() { + return uploadConcurrency; + } + + public void setUploadConcurrency(int uploadConcurrency) { + this.uploadConcurrency = uploadConcurrency; + } + + public int getMaxRetries() { + return maxRetries; + } + + public void setMaxRetries(int maxRetries) { + this.maxRetries = maxRetries; + } + + public long getRetryIntervalMillis() { + return retryIntervalMillis; + } + + public void setRetryIntervalMillis(long retryIntervalMillis) { + this.retryIntervalMillis = retryIntervalMillis; + } + + public ProgressListener getProgressListener() { + return progressListener; + } + + public void setProgressListener(ProgressListener progressListener) { + this.progressListener = progressListener; + } + + public long getPartSizeBytes() { + return partSizeBytes; + } + + public void setPartSizeBytes(long partSizeBytes) { + this.partSizeBytes = partSizeBytes; + } + @Override public String toString() { return "UploadFilesRequest{" + "sourceFilePath='" + sourceFilePath + '\'' + ", targetVolumePath='" + targetVolumePath + '\'' + + ", uploadConcurrency=" + uploadConcurrency + + ", maxRetries=" + maxRetries + + ", retryIntervalMillis=" + retryIntervalMillis + + ", progressListener=" + (progressListener != null) + + ", partSizeBytes=" + partSizeBytes + '}'; } @@ -78,10 +155,20 @@ public static UploadFilesRequestBuilder builder() { public static class UploadFilesRequestBuilder { private String sourceFilePath; private String targetVolumePath; + private int uploadConcurrency; + private int maxRetries; + private long retryIntervalMillis; + private ProgressListener progressListener; + private long partSizeBytes; private UploadFilesRequestBuilder() { this.sourceFilePath = ""; this.targetVolumePath = ""; + this.uploadConcurrency = 5; + this.maxRetries = 5; + this.retryIntervalMillis = 5000L; + this.progressListener = null; + this.partSizeBytes = 0L; } public UploadFilesRequestBuilder sourceFilePath(String sourceFilePath) { @@ -94,8 +181,38 @@ public UploadFilesRequestBuilder targetVolumePath(String targetVolumePath) { return this; } + public UploadFilesRequestBuilder uploadConcurrency(int uploadConcurrency) { + this.uploadConcurrency = uploadConcurrency; + return this; + } + + public UploadFilesRequestBuilder maxRetries(int maxRetries) { + this.maxRetries = maxRetries; + return this; + } + + public UploadFilesRequestBuilder retryIntervalMillis(long retryIntervalMillis) { + this.retryIntervalMillis = retryIntervalMillis; + return this; + } + + public UploadFilesRequestBuilder progressListener(ProgressListener progressListener) { + this.progressListener = progressListener; + return this; + } + + public UploadFilesRequestBuilder partSizeBytes(long partSizeBytes) { + this.partSizeBytes = partSizeBytes; + return this; + } + public UploadFilesRequest build() { return new UploadFilesRequest(this); } } + + @FunctionalInterface + public interface ProgressListener { + void onProgress(UploadProgress progress); + } } diff --git a/sdk-bulkwriter/src/main/java/io/milvus/bulkwriter/storage/StorageClient.java b/sdk-bulkwriter/src/main/java/io/milvus/bulkwriter/storage/StorageClient.java index fdb770155..1074fb722 100644 --- a/sdk-bulkwriter/src/main/java/io/milvus/bulkwriter/storage/StorageClient.java +++ b/sdk-bulkwriter/src/main/java/io/milvus/bulkwriter/storage/StorageClient.java @@ -28,4 +28,22 @@ public interface StorageClient { boolean checkBucketExist(String bucketName) throws Exception; void putObject(File file, String bucketName, String objectKey) throws Exception; + + default void putObject(File file, String bucketName, String objectKey, + UploadProgressListener progressListener) throws Exception { + putObject(file, bucketName, objectKey); + } + + default void putObject(File file, String bucketName, String objectKey, + UploadProgressListener progressListener, long partSizeBytes) throws Exception { + putObject(file, bucketName, objectKey, progressListener); + } + + default void close() { + } + + @FunctionalInterface + interface UploadProgressListener { + void onProgress(long bytes); + } } diff --git a/sdk-bulkwriter/src/main/java/io/milvus/bulkwriter/storage/client/MinioStorageClient.java b/sdk-bulkwriter/src/main/java/io/milvus/bulkwriter/storage/client/MinioStorageClient.java index 4e7688c8d..f15be40ab 100644 --- a/sdk-bulkwriter/src/main/java/io/milvus/bulkwriter/storage/client/MinioStorageClient.java +++ b/sdk-bulkwriter/src/main/java/io/milvus/bulkwriter/storage/client/MinioStorageClient.java @@ -23,6 +23,7 @@ import io.milvus.bulkwriter.common.clientenum.CloudStorage; import io.milvus.bulkwriter.model.CompleteMultipartUploadOutputModel; import io.milvus.bulkwriter.storage.StorageClient; +import io.milvus.exception.ParamException; import io.minio.BucketExistsArgs; import io.minio.MinioAsyncClient; import io.minio.ObjectWriteResponse; @@ -50,7 +51,9 @@ import java.io.File; import java.io.FileInputStream; +import java.io.FilterInputStream; import java.io.IOException; +import java.io.InputStream; import java.security.InvalidKeyException; import java.security.NoSuchAlgorithmException; import java.util.concurrent.CompletableFuture; @@ -61,10 +64,14 @@ public class MinioStorageClient extends MinioAsyncClient implements StorageClient { private static final Logger logger = LoggerFactory.getLogger(MinioStorageClient.class); private static final String UPLOAD_ID = "uploadId"; + private static final long MIN_MULTIPART_PART_SIZE = 5L * MB; + private static final long TARGET_MULTIPART_PART_COUNT = 1000L; + private static final long MAX_MULTIPART_PART_COUNT = 10000L; + private final boolean closeHttpClient; - - protected MinioStorageClient(MinioAsyncClient client) { + protected MinioStorageClient(MinioAsyncClient client, boolean closeHttpClient) { super(client); + this.closeHttpClient = closeHttpClient; } public static MinioStorageClient getStorageClient(String cloudName, @@ -74,6 +81,7 @@ public static MinioStorageClient getStorageClient(String cloudName, String sessionToken, String region, OkHttpClient httpClient) { + boolean closeHttpClient = httpClient == null; MinioAsyncClient.Builder minioClientBuilder = MinioAsyncClient.builder() .endpoint(endpoint); @@ -97,7 +105,7 @@ public static MinioStorageClient getStorageClient(String cloudName, minioClient.enableVirtualStyleEndpoint(); } - return new MinioStorageClient(minioClient); + return new MinioStorageClient(minioClient, closeHttpClient); } private static OkHttpClient buildAuthorizedClient(OkHttpClient httpClient, String sessionToken) { @@ -130,14 +138,64 @@ public Long getObjectEntity(String bucketName, String objectKey) throws Exceptio } public void putObject(File file, String bucketName, String objectKey) throws Exception { + putObject(file, bucketName, objectKey, null, 0L); + } + + @Override + public void putObject(File file, String bucketName, String objectKey, + UploadProgressListener progressListener) throws Exception { + putObject(file, bucketName, objectKey, progressListener, 0L); + } + + @Override + public void putObject(File file, String bucketName, String objectKey, + UploadProgressListener progressListener, long partSizeBytes) throws Exception { logger.info("uploading file, fileName:{}, size:{} bytes", file.getAbsolutePath(), file.length()); - FileInputStream fileInputStream = new FileInputStream(file); - PutObjectArgs putObjectArgs = PutObjectArgs.builder() - .bucket(bucketName) - .object(objectKey) - .stream(fileInputStream, file.length(), 5 * MB) - .build(); - putObject(putObjectArgs).get(); + long uploadPartSize = calculateUploadPartSize(file.length(), partSizeBytes); + try (InputStream fileInputStream = new ProgressInputStream(new FileInputStream(file), progressListener)) { + PutObjectArgs putObjectArgs = PutObjectArgs.builder() + .bucket(bucketName) + .object(objectKey) + .stream(fileInputStream, file.length(), uploadPartSize) + .build(); + putObject(putObjectArgs).get(); + } + } + + @Override + public void close() { + if (!closeHttpClient || httpClient == null) { + return; + } + httpClient.dispatcher().executorService().shutdown(); + httpClient.connectionPool().evictAll(); + if (httpClient.cache() != null) { + try { + httpClient.cache().close(); + } catch (IOException e) { + logger.warn("Failed to close MinIO HTTP client cache", e); + } + } + } + + static long calculateUploadPartSize(long fileSize, long requestedPartSizeBytes) { + if (requestedPartSizeBytes > 0) { + if (requestedPartSizeBytes < MIN_MULTIPART_PART_SIZE) { + throw new ParamException("partSizeBytes must be at least " + MIN_MULTIPART_PART_SIZE + " bytes"); + } + return requestedPartSizeBytes; + } + if (fileSize <= 0) { + return MIN_MULTIPART_PART_SIZE; + } + long targetPartSize = divideCeil(fileSize, TARGET_MULTIPART_PART_COUNT); + long maxPartCountSize = divideCeil(fileSize, MAX_MULTIPART_PART_COUNT); + long partSize = Math.max(MIN_MULTIPART_PART_SIZE, Math.max(targetPartSize, maxPartCountSize)); + return divideCeil(partSize, MB) * MB; + } + + private static long divideCeil(long value, long divisor) { + return (value + divisor - 1) / divisor; } public boolean checkBucketExist(String bucketName) throws Exception { @@ -225,3 +283,36 @@ protected CompletableFuture completeMultipartUploadAsync(St }); } } + +class ProgressInputStream extends FilterInputStream { + private final StorageClient.UploadProgressListener progressListener; + + ProgressInputStream(InputStream inputStream, StorageClient.UploadProgressListener progressListener) { + super(inputStream); + this.progressListener = progressListener; + } + + @Override + public int read() throws IOException { + int value = super.read(); + if (value != -1) { + notifyProgress(1); + } + return value; + } + + @Override + public int read(byte[] bytes, int offset, int length) throws IOException { + int count = super.read(bytes, offset, length); + if (count > 0) { + notifyProgress(count); + } + return count; + } + + private void notifyProgress(long bytes) { + if (progressListener != null) { + progressListener.onProgress(bytes); + } + } +} diff --git a/sdk-bulkwriter/src/test/java/io/milvus/bulkwriter/VolumeManagerTest.java b/sdk-bulkwriter/src/test/java/io/milvus/bulkwriter/VolumeManagerTest.java index bacf7cece..b2cae3fde 100644 --- a/sdk-bulkwriter/src/test/java/io/milvus/bulkwriter/VolumeManagerTest.java +++ b/sdk-bulkwriter/src/test/java/io/milvus/bulkwriter/VolumeManagerTest.java @@ -20,15 +20,34 @@ package io.milvus.bulkwriter; import com.google.gson.Gson; +import io.milvus.bulkwriter.model.UploadProgress; +import io.milvus.bulkwriter.request.volume.ApplyVolumeRequest; import io.milvus.bulkwriter.request.volume.CreateVolumeRequest; import io.milvus.bulkwriter.request.volume.DeleteVolumeRequest; import io.milvus.bulkwriter.request.volume.DescribeVolumeRequest; import io.milvus.bulkwriter.request.volume.ListVolumesRequest; +import io.milvus.bulkwriter.request.volume.UploadFilesRequest; +import io.milvus.bulkwriter.response.ApplyVolumeResponse; import io.milvus.bulkwriter.response.volume.ListVolumesResponse; import io.milvus.bulkwriter.response.volume.VolumeInfo; +import io.milvus.bulkwriter.storage.StorageClient; +import io.milvus.exception.ParamException; import org.junit.jupiter.api.Test; +import java.io.File; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.CyclicBarrier; +import java.util.concurrent.atomic.AtomicInteger; import static org.junit.jupiter.api.Assertions.*; @@ -36,6 +55,216 @@ public class VolumeManagerTest { private final Gson gson = new Gson(); + @Test + public void testUploadFilesRequestDefaultUploadOptions() { + UploadFilesRequest request = UploadFilesRequest.builder() + .sourceFilePath("/tmp/data") + .targetVolumePath("data/") + .build(); + + assertEquals(5, request.getUploadConcurrency()); + assertEquals(5, request.getMaxRetries()); + assertEquals(5000L, request.getRetryIntervalMillis()); + assertNull(request.getProgressListener()); + assertEquals(0L, request.getPartSizeBytes()); + } + + @Test + public void testUploadFilesRequestCustomUploadOptions() { + UploadFilesRequest.ProgressListener progressListener = progress -> { + }; + UploadFilesRequest request = UploadFilesRequest.builder() + .sourceFilePath("/tmp/data") + .targetVolumePath("data/") + .uploadConcurrency(3) + .maxRetries(2) + .retryIntervalMillis(100L) + .progressListener(progressListener) + .partSizeBytes(16L * 1024L * 1024L) + .build(); + + assertEquals(3, request.getUploadConcurrency()); + assertEquals(2, request.getMaxRetries()); + assertEquals(100L, request.getRetryIntervalMillis()); + assertSame(progressListener, request.getProgressListener()); + assertEquals(16L * 1024L * 1024L, request.getPartSizeBytes()); + } + + @Test + public void testConcurrentUploadFilesUseRequestLocalVolumeState() throws Exception { + Path fileA = Files.createTempFile("volume-upload-a", ".txt"); + Path fileB = Files.createTempFile("volume-upload-b", ".txt"); + Files.write(fileA, Collections.singletonList("a")); + Files.write(fileB, Collections.singletonList("b")); + CyclicBarrier barrier = new CyclicBarrier(2); + ConcurrentVolumeFileManager manager = new ConcurrentVolumeFileManager(barrier); + + try { + CompletableFuture futureA = manager.uploadFilesAsync(UploadFilesRequest.builder() + .sourceFilePath(fileA.toString()) + .targetVolumePath("a/") + .uploadConcurrency(1) + .build()); + CompletableFuture futureB = manager.uploadFilesAsync(UploadFilesRequest.builder() + .sourceFilePath(fileB.toString()) + .targetVolumePath("b/") + .uploadConcurrency(1) + .build()); + + futureA.get(); + futureB.get(); + + assertEquals(Collections.singletonList("prefix-a/a/" + fileA.getFileName()), + manager.uploadsByBucket.get("bucket-a")); + assertEquals(Collections.singletonList("prefix-b/b/" + fileB.getFileName()), + manager.uploadsByBucket.get("bucket-b")); + } finally { + Files.deleteIfExists(fileA); + Files.deleteIfExists(fileB); + } + } + + @Test + public void testUploadFilesProgressListenerReceivesProgress() throws Exception { + Path file = Files.createTempFile("volume-upload-progress", ".txt"); + Files.write(file, Collections.singletonList("data")); + ConcurrentVolumeFileManager manager = new ConcurrentVolumeFileManager(new CyclicBarrier(1)); + List progressEvents = new CopyOnWriteArrayList<>(); + + try { + manager.uploadFiles(UploadFilesRequest.builder() + .sourceFilePath(file.toString()) + .targetVolumePath("progress/") + .uploadConcurrency(1) + .progressListener(progressEvents::add) + .build()); + + assertEquals(3, progressEvents.size()); + assertEquals(100.0, progressEvents.get(0).getPercent()); + assertEquals(file.toString(), progressEvents.get(0).getCurrentFile()); + assertEquals("", progressEvents.get(progressEvents.size() - 1).getCurrentFile()); + } finally { + Files.deleteIfExists(file); + } + } + + @Test + public void testUploadFilesProgressListenerFailureDoesNotRetry() throws Exception { + Path file = Files.createTempFile("volume-upload-progress-failure", ".txt"); + Files.write(file, Collections.singletonList("data")); + ConcurrentVolumeFileManager manager = new ConcurrentVolumeFileManager(new CyclicBarrier(1)); + + try { + UploadFilesRequest request = UploadFilesRequest.builder() + .sourceFilePath(file.toString()) + .targetVolumePath("failure/") + .uploadConcurrency(1) + .maxRetries(3) + .progressListener(progress -> { + throw new RuntimeException("stop"); + }) + .build(); + + assertThrows(Exception.class, () -> manager.uploadFiles(request)); + assertEquals(1, manager.uploadsByBucket.get("bucket-failure").size()); + } finally { + Files.deleteIfExists(file); + } + } + + @Test + public void testUploadFilesClosesRequestStorageClient() throws Exception { + Path file = Files.createTempFile("volume-upload-close", ".txt"); + Files.write(file, Collections.singletonList("data")); + ConcurrentVolumeFileManager manager = new ConcurrentVolumeFileManager(new CyclicBarrier(1)); + + try { + manager.uploadFiles(UploadFilesRequest.builder() + .sourceFilePath(file.toString()) + .targetVolumePath("close/") + .uploadConcurrency(1) + .build()); + + assertEquals(Collections.singletonList("bucket-close"), manager.closedBuckets); + } finally { + Files.deleteIfExists(file); + } + } + + @Test + public void testUploadFilesMaxRetriesZeroDisablesRetry() throws Exception { + Path file = Files.createTempFile("volume-upload-no-retry", ".txt"); + Files.write(file, Collections.singletonList("data")); + ConcurrentVolumeFileManager manager = new ConcurrentVolumeFileManager(new CyclicBarrier(1)); + manager.uploadException = new IOException("temporary"); + + try { + UploadFilesRequest request = UploadFilesRequest.builder() + .sourceFilePath(file.toString()) + .targetVolumePath("no-retry/") + .uploadConcurrency(1) + .maxRetries(0) + .retryIntervalMillis(0L) + .build(); + + assertThrows(Exception.class, () -> manager.uploadFiles(request)); + assertEquals(1, manager.attemptsByBucket.get("bucket-no-retry").get()); + } finally { + Files.deleteIfExists(file); + } + } + + @Test + public void testUploadFilesDoesNotRetryNonRetryableFailure() throws Exception { + Path file = Files.createTempFile("volume-upload-param-failure", ".txt"); + Files.write(file, Collections.singletonList("data")); + ConcurrentVolumeFileManager manager = new ConcurrentVolumeFileManager(new CyclicBarrier(1)); + manager.uploadException = new ParamException("invalid part size"); + + try { + UploadFilesRequest request = UploadFilesRequest.builder() + .sourceFilePath(file.toString()) + .targetVolumePath("param-failure/") + .uploadConcurrency(1) + .maxRetries(3) + .retryIntervalMillis(0L) + .build(); + + assertThrows(Exception.class, () -> manager.uploadFiles(request)); + assertEquals(1, manager.attemptsByBucket.get("bucket-param-failure").get()); + } finally { + Files.deleteIfExists(file); + } + } + + @Test + public void testUploadFilesProgressUsesCanonicalPathForRelativeSingleFile() throws Exception { + Path dir = Paths.get("target", "volume-relative-progress"); + Files.createDirectories(dir); + Path file = dir.resolve("relative-progress.txt"); + Files.write(file, Collections.singletonList("data")); + String relativeFilePath = Paths.get("").toAbsolutePath().relativize(file.toAbsolutePath()).toString(); + ConcurrentVolumeFileManager manager = new ConcurrentVolumeFileManager(new CyclicBarrier(1)); + List progressEvents = new CopyOnWriteArrayList<>(); + + try { + manager.uploadFiles(UploadFilesRequest.builder() + .sourceFilePath(relativeFilePath) + .targetVolumePath("relative-progress/") + .uploadConcurrency(1) + .progressListener(progressEvents::add) + .build()); + + assertFalse(progressEvents.isEmpty()); + long fileSize = file.toFile().length(); + assertTrue(progressEvents.stream().allMatch(progress -> progress.getUploadedBytes() <= fileSize)); + assertEquals(fileSize, progressEvents.get(progressEvents.size() - 1).getUploadedBytes()); + } finally { + Files.deleteIfExists(file); + Files.deleteIfExists(dir); + } + } + // ========== CreateVolumeRequest Tests ========== @Test @@ -446,4 +675,117 @@ public void testListVolumesResponseDeserialization() { assertEquals("MANAGED", response.getVolumes().get(0).getType()); assertEquals("EXTERNAL", response.getVolumes().get(1).getType()); } + + private static class ConcurrentVolumeFileManager extends VolumeFileManager { + private final Gson gson = new Gson(); + private final CyclicBarrier barrier; + private final Map> uploadsByBucket = new ConcurrentHashMap<>(); + private final Map attemptsByBucket = new ConcurrentHashMap<>(); + private final List closedBuckets = new CopyOnWriteArrayList<>(); + private volatile Exception uploadException; + + private ConcurrentVolumeFileManager(CyclicBarrier barrier) throws Exception { + super(VolumeFileManagerParam.newBuilder() + .withCloudEndpoint("https://api.cloud.zilliz.com") + .withApiKey("test-api-key") + .withVolumeName("test-volume") + .build()); + this.barrier = barrier; + } + + @Override + protected String applyVolume(ApplyVolumeRequest applyVolumeRequest) { + String name = applyVolumeRequest.getPath().replace("/", ""); + ApplyVolumeResponse response = ApplyVolumeResponse.builder() + .endpoint("s3.amazonaws.com") + .cloud("aws") + .region("us-west-2") + .bucketName("bucket-" + name) + .uploadPath("") + .credentials(ApplyVolumeResponse.Credentials.builder() + .tmpAK("ak") + .tmpSK("sk") + .sessionToken("token") + .expireTime("2099-12-31T23:59:59Z") + .build()) + .condition(ApplyVolumeResponse.Condition.builder() + .maxContentLength(1024L * 1024L) + .build()) + .volumeName("test-volume") + .volumePrefix("prefix-" + name + "/") + .build(); + return gson.toJson(response); + } + + @Override + protected StorageClient createStorageClient(ApplyVolumeResponse applyVolumeResponse) { + uploadsByBucket.putIfAbsent(applyVolumeResponse.getBucketName(), new CopyOnWriteArrayList<>()); + attemptsByBucket.putIfAbsent(applyVolumeResponse.getBucketName(), new AtomicInteger(0)); + return new RecordingStorageClient( + barrier, applyVolumeResponse.getBucketName(), uploadsByBucket, attemptsByBucket, + closedBuckets, () -> uploadException); + } + } + + @FunctionalInterface + private interface UploadExceptionSupplier { + Exception get(); + } + + private static class RecordingStorageClient implements StorageClient { + private final CyclicBarrier barrier; + private final String bucketName; + private final Map> uploadsByBucket; + private final Map attemptsByBucket; + private final List closedBuckets; + private final UploadExceptionSupplier uploadExceptionSupplier; + + private RecordingStorageClient(CyclicBarrier barrier, String bucketName, + Map> uploadsByBucket, + Map attemptsByBucket, + List closedBuckets, + UploadExceptionSupplier uploadExceptionSupplier) { + this.barrier = barrier; + this.bucketName = bucketName; + this.uploadsByBucket = uploadsByBucket; + this.attemptsByBucket = attemptsByBucket; + this.closedBuckets = closedBuckets; + this.uploadExceptionSupplier = uploadExceptionSupplier; + } + + @Override + public Long getObjectEntity(String bucketName, String objectKey) { + return 0L; + } + + @Override + public boolean checkBucketExist(String bucketName) { + return true; + } + + @Override + public void putObject(File file, String bucketName, String objectKey) throws Exception { + putObject(file, bucketName, objectKey, null); + } + + @Override + public void putObject(File file, String bucketName, String objectKey, + UploadProgressListener progressListener) throws Exception { + barrier.await(); + attemptsByBucket.get(bucketName).incrementAndGet(); + Exception uploadException = uploadExceptionSupplier.get(); + if (uploadException != null) { + throw uploadException; + } + uploadsByBucket.get(bucketName).add(objectKey); + if (progressListener != null) { + progressListener.onProgress(file.length()); + } + } + + @Override + public void close() { + closedBuckets.add(bucketName); + } + } } diff --git a/sdk-bulkwriter/src/test/java/io/milvus/bulkwriter/storage/client/ProgressInputStreamTest.java b/sdk-bulkwriter/src/test/java/io/milvus/bulkwriter/storage/client/ProgressInputStreamTest.java new file mode 100644 index 000000000..321ee144e --- /dev/null +++ b/sdk-bulkwriter/src/test/java/io/milvus/bulkwriter/storage/client/ProgressInputStreamTest.java @@ -0,0 +1,68 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package io.milvus.bulkwriter.storage.client; + +import io.milvus.bulkwriter.storage.StorageClient; +import io.milvus.exception.ParamException; +import org.junit.jupiter.api.Test; + +import java.io.ByteArrayInputStream; +import java.util.concurrent.atomic.AtomicLong; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +public class ProgressInputStreamTest { + + @Test + public void testProgressCallbackCountsReadBytes() throws Exception { + byte[] bytes = new byte[]{1, 2, 3, 4, 5}; + AtomicLong uploadedBytes = new AtomicLong(0); + StorageClient.UploadProgressListener listener = uploadedBytes::addAndGet; + + try (ProgressInputStream inputStream = new ProgressInputStream(new ByteArrayInputStream(bytes), listener)) { + byte[] buffer = new byte[2]; + assertEquals(2, inputStream.read(buffer)); + assertEquals(2, uploadedBytes.get()); + assertEquals(2, inputStream.read(buffer)); + assertEquals(4, uploadedBytes.get()); + assertEquals(5, inputStream.read()); + assertEquals(5, uploadedBytes.get()); + assertEquals(-1, inputStream.read(buffer)); + assertEquals(5, uploadedBytes.get()); + } + } + + @Test + public void testCalculateUploadPartSizeAutoAndExplicit() { + assertEquals(5L * 1024L * 1024L, MinioStorageClient.calculateUploadPartSize(1L, 0L)); + assertEquals(7L * 1024L * 1024L, + MinioStorageClient.calculateUploadPartSize(1L, 7L * 1024L * 1024L)); + long partSize = MinioStorageClient.calculateUploadPartSize(20L * 1024L * 1024L * 1024L, 0L); + assertEquals(0L, partSize % (1024L * 1024L)); + assertTrue(partSize > 5L * 1024L * 1024L); + } + + @Test + public void testCalculateUploadPartSizeRejectsTooSmallExplicitValue() { + assertThrows(ParamException.class, () -> MinioStorageClient.calculateUploadPartSize(1L, 1L)); + } +}