Examples with UploadFileChunkRequest ai.nightfall.scan.model.UploadFileChunkRequest used on opensource projects

Example 1 with UploadFileChunkRequest

use of ai.nightfall.scan.model.UploadFileChunkRequest in project nightfall-java-sdk by nightfallai.

the class NightfallClient method doChunkedUpload.

private boolean doChunkedUpload(FileUpload upload, InputStream content, Instant deadline, AtomicReference<BaseNightfallException> uploadException) {
    // Use a semaphore to avoid loading the entire stream into memory
    int numPermits = this.fileUploadConcurrency;
    Semaphore semaphore = new Semaphore(numPermits);
    AtomicBoolean allChunksSucceed = new AtomicBoolean(true);
    for (int offset = 0; offset < upload.getFileSizeBytes(); offset += upload.getChunkSize()) {
        semaphore.acquireUninterruptibly();
        checkFileUploadDeadline(deadline);
        if (!allChunksSucceed.get()) {
            return false;
        }
        UploadFileChunkRequest chunkReq = new UploadFileChunkRequest(upload.getFileID(), offset);
        byte[] data = new byte[(int) upload.getChunkSize()];
        try {
            int bytesRead = content.read(data);
            boolean notLastChunk = offset + upload.getChunkSize() < upload.getFileSizeBytes();
            if (bytesRead < data.length && notLastChunk) {
                semaphore.release();
                throw new NightfallClientException("failed to read data from input stream");
            } else if (bytesRead < data.length) {
                data = Arrays.copyOfRange(data, 0, bytesRead);
            }
        } catch (IOException e) {
            semaphore.release();
            throw new NightfallClientException("reading content to upload: " + e.getMessage());
        }
        chunkReq.setContent(data);
        this.executor.execute(() -> {
            try {
                this.uploadFileChunk(chunkReq);
            } catch (BaseNightfallException e) {
                allChunksSucceed.set(false);
                uploadException.set(e);
            } catch (Throwable t) {
                allChunksSucceed.set(false);
            } finally {
                semaphore.release();
            }
        });
    }
    while (true) {
        try {
            // Attempt to acquire all permits; this is only possible when all chunks have been uploaded.
            // Allow spurious wake-ups in case the caller puts a deadline on the operation.
            boolean success = semaphore.tryAcquire(numPermits, wakeupDurationMillis, TimeUnit.MILLISECONDS);
            if (success) {
                return allChunksSucceed.get();
            }
            checkFileUploadDeadline(deadline);
        } catch (InterruptedException e) {
            throw new NightfallClientException("interrupted while waiting for upload to complete");
        }
    }
}