fix: persist partial bytes and update offset on interrupted input stream across all storage backends (#117)

* fix(ci): update maven-dependency-submission-action to v5.0.0

* fix: Code duplication

* fix: persist partial bytes and update offset on interrupted input stream across all storage backends

Author: tomdesair

src/main/java/me/desair/tus/server/upload/UploadLockingService.java

@@ -43,7 +43,7 @@ public interface UploadLockingService {
 
   /**
    * Register the input stream associated with the active request URI so that it can be interrupted
-   * if lock contention occurs.
+   * if lock contention occurs or during service shutdown.
    *
    * @param requestUri The request URI of the active request
    * @param inputStream The input stream of the active request
@@ -53,8 +53,9 @@ default void registerInputStream(String requestUri, java.io.InputStream inputStr
   }
 
   /**
-   * Request that the lock for the given request URI be released. This might involve interrupting
-   * the active request's input stream.
+   * Request that the lock for the given request URI be released. This interrupts the active
+   * request's input stream locally or signals other cluster replicas to interrupt their active
+   * stream via a stop signal.
    *
    * @param requestUri The request URI of the upload lock to release
    */
@@ -63,8 +64,8 @@ default void requestLockRelease(String requestUri) {
   }
 
   /**
-   * Closes resources and shuts down any background watchdog threads associated with this locking
-   * service.
+   * Closes resources, interrupts any active in-flight request streams, and shuts down background
+   * watchdog threads associated with this locking service.
    *
    * @throws IOException If closing fails
    */

src/main/java/me/desair/tus/server/upload/UploadStorageService.java

@@ -37,13 +37,20 @@ public interface UploadStorageService {
   String getUploadUri();
 
   /**
-   * Append the bytes in the give {@link InputStream} to the upload with the given ID starting at
+   * Append the bytes in the given {@link InputStream} to the upload with the given ID starting at
    * the provided offset. This method also updates the {@link UploadInfo} corresponding to this
    * upload. The Upload Storage server should not exceed its max upload size when writing bytes.
    *
+   * <p>If the input stream is interrupted or encounters an {@link IOException} during reading (e.g.
+   * network disconnect, socket closed, JVM shutdown signal, or lock release interruption), the
+   * storage service MUST cleanly persist all bytes received up to the interruption point, update
+   * the {@link UploadInfo} offset and metadata accordingly, and re-throw the exception.
+   *
    * @param upload The ID of the upload
    * @param inputStream The input stream containing the bytes to append
    * @return The new {@link UploadInfo} for this upload
+   * @throws IOException If saving bytes or reading from the input stream fails
+   * @throws TusException If upload validation fails or storage constraints are violated
    */
   UploadInfo append(UploadInfo upload, InputStream inputStream) throws IOException, TusException;
 

src/main/java/me/desair/tus/server/upload/azure/AzureBlobLockingService.java

@@ -96,6 +96,11 @@ protected void cleanupOnClose() throws IOException {
       Utils.interruptThread(watchdogThread);
       watchdogThread = null;
     }
+    for (WeakReference<InterruptibleInputStream> streamRef : activeStreams.values()) {
+      if (streamRef != null) {
+        Utils.interruptStream(streamRef.get());
+      }
+    }
     activeStreams.clear();
   }
 

src/main/java/me/desair/tus/server/upload/azure/AzureBlobStorageService.java

@@ -210,12 +210,15 @@ public UploadInfo append(UploadInfo upload, InputStream inputStream)
 
     long totalAppended = 0L;
     boolean streamFinished = false;
+    IOException streamException = null;
 
     File firstChunkFile = File.createTempFile("tus-azure-chunk-", ".tmp", tempBufferDir.toFile());
     try {
       // 4. Read first chunk from incoming payload stream into local disk buffer
-      long firstChunkSize = readChunk(inputStream, firstChunkFile, optimalBlockSize);
+      ReadChunkResult firstChunkResult = readChunk(inputStream, firstChunkFile, optimalBlockSize);
+      long firstChunkSize = firstChunkResult.bytesRead;
       totalAppended += firstChunkSize;
+      streamException = firstChunkResult.exception;
 
       validateMaxAppendSize(totalAppended, effectiveMaxAppendSize);
 
@@ -227,17 +230,25 @@ public UploadInfo append(UploadInfo upload, InputStream inputStream)
         streamFinished = true;
       }
 
-      if (totalBuffered < optimalBlockSize && !isUploadComplete && streamFinished) {
+      // If the cumulative buffered data is below optimalBlockSize and the upload is not yet
+      // finished, and the stream has reached EOF or was interrupted by an IOException:
+      // Keep data in the temporary .part blob instead of committing a sub-optimal block to the
+      // Block Blob.
+      if (totalBuffered < optimalBlockSize
+          && !isUploadComplete
+          && (streamFinished || streamException != null)) {
         // Small append under block size threshold: buffer data to .part blob directly
-        bufferToPartBlob(partBlob, existingPartSize, firstChunkFile, firstChunkSize);
+        if (firstChunkSize > 0) {
+          bufferToPartBlob(partBlob, existingPartSize, firstChunkFile, firstChunkSize);
+        }
       } else {
         // Data exceeds block size threshold: stage blocks to Azure Block Blob
         stagePartBlobIfPresent(partBlob, existingPartSize, blockBlobClient, blockIds);
         stageChunkFile(firstChunkFile, firstChunkSize, blockBlobClient, blockIds);
 
         // Process any remaining chunks from input stream
-        if (!streamFinished) {
-          totalAppended +=
+        if (!streamFinished && streamException == null) {
+          ProcessChunksResult remainingResult =
               processRemainingChunks(
                   inputStream,
                   optimalBlockSize,
@@ -247,13 +258,21 @@ public UploadInfo append(UploadInfo upload, InputStream inputStream)
                   blockBlobClient,
                   blockIds,
                   totalAppended);
+          totalAppended += remainingResult.additionalAppended;
+          if (remainingResult.exception != null) {
+            streamException = remainingResult.exception;
+          }
         }
 
         // Commit updated block ID list on Azure so staged blocks become committed and readable
-        blockBlobClient.commitBlockList(blockIds, true);
+        if (!blockIds.isEmpty()) {
+          blockBlobClient.commitBlockList(blockIds, true);
+        }
       }
 
-      validateMinAppendSize(totalAppended);
+      if (streamException == null) {
+        validateMinAppendSize(totalAppended);
+      }
 
       // 5. Update UploadInfo offset, expiration timestamp, and optional deduplication state
       upload.setOffset(upload.getOffset() + totalAppended);
@@ -268,6 +287,11 @@ public UploadInfo append(UploadInfo upload, InputStream inputStream)
       }
 
       saveUploadInfo(upload);
+
+      if (streamException != null) {
+        throw streamException;
+      }
+
       return upload;
     } finally {
       deleteFileQuietly(firstChunkFile);
@@ -702,10 +726,31 @@ private long getPartBlobSize(BlobClient partBlob) {
     }
   }
 
-  /** Reads up to maxBytes from InputStream into target File. */
-  private long readChunk(InputStream is, File targetFile, long maxBytes) throws IOException {
+  private static class ReadChunkResult {
+    final long bytesRead;
+    final IOException exception;
+
+    ReadChunkResult(long bytesRead, IOException exception) {
+      this.bytesRead = bytesRead;
+      this.exception = exception;
+    }
+  }
+
+  private static class ProcessChunksResult {
+    final long additionalAppended;
+    final IOException exception;
+
+    ProcessChunksResult(long additionalAppended, IOException exception) {
+      this.additionalAppended = additionalAppended;
+      this.exception = exception;
+    }
+  }
+
+  /** Reads up to maxBytes from InputStream into target File, preserving read count on exception. */
+  private ReadChunkResult readChunk(InputStream is, File targetFile, long maxBytes) {
     long totalRead = 0L;
     byte[] buffer = new byte[8192];
+    IOException readException = null;
     try (FileOutputStream fos = new FileOutputStream(targetFile)) {
       while (totalRead < maxBytes) {
         int lenToRead = (int) Math.min(buffer.length, maxBytes - totalRead);
@@ -716,8 +761,10 @@ private long readChunk(InputStream is, File targetFile, long maxBytes) throws IO
         fos.write(buffer, 0, read);
         totalRead += read;
       }
+    } catch (IOException e) {
+      readException = e;
     }
-    return totalRead;
+    return new ReadChunkResult(totalRead, readException);
   }
 
   /** Retrieves committed block IDs from Azure Block Blob. */
@@ -786,8 +833,8 @@ private void stageChunkFile(
     }
   }
 
-  /** Processes remaining payload chunks from stream until EOF. */
-  private long processRemainingChunks(
+  /** Processes remaining payload chunks from stream until EOF or interruption. */
+  private ProcessChunksResult processRemainingChunks(
       InputStream inputStream,
       long optimalBlockSize,
       Long effectiveMaxAppendSize,
@@ -796,14 +843,19 @@ private long processRemainingChunks(
       BlockBlobClient blockBlobClient,
       List<String> blockIds,
       long currentTotalAppended)
-      throws IOException, MaxAppendSizeExceededException {
+      throws MaxAppendSizeExceededException, IOException {
     long additionalAppended = 0L;
     boolean streamFinished = false;
+    IOException exception = null;
 
     while (!streamFinished) {
-      File chunkFile = File.createTempFile("tus-azure-chunk-", ".tmp", tempBufferDir.toFile());
+      File chunkFile = null;
       try {
-        long chunkSize = readChunk(inputStream, chunkFile, optimalBlockSize);
+        chunkFile = File.createTempFile("tus-azure-chunk-", ".tmp", tempBufferDir.toFile());
+        ReadChunkResult chunkResult = readChunk(inputStream, chunkFile, optimalBlockSize);
+        long chunkSize = chunkResult.bytesRead;
+        exception = chunkResult.exception;
+
         if (chunkSize <= 0) {
           break;
         }
@@ -820,11 +872,17 @@ private long processRemainingChunks(
         } else {
           stageChunkFile(chunkFile, chunkSize, blockBlobClient, blockIds);
         }
+
+        if (exception != null) {
+          break;
+        }
       } finally {
-        deleteFileQuietly(chunkFile);
+        if (chunkFile != null) {
+          deleteFileQuietly(chunkFile);
+        }
       }
     }
-    return additionalAppended;
+    return new ProcessChunksResult(additionalAppended, exception);
   }
 
   /** Buffers incoming data to temporary .part blob when under block threshold. */

src/main/java/me/desair/tus/server/upload/disk/DiskLockingService.java

@@ -61,6 +61,11 @@ protected void cleanupOnClose() throws IOException {
       Utils.interruptThread(watchdogThread);
       watchdogThread = null;
     }
+    for (WeakReference<InterruptibleInputStream> streamRef : activeLocks.values()) {
+      if (streamRef != null) {
+        Utils.interruptStream(streamRef.get());
+      }
+    }
     activeLocks.clear();
   }
 

src/main/java/me/desair/tus/server/upload/s3/S3LockingService.java

@@ -313,6 +313,9 @@ private void writeStopSignal(UploadId uploadId) {
   @Override
   protected void cleanupOnClose() throws IOException {
     Utils.shutdownExecutor(watchdogExecutor);
+    for (InputStream stream : activeInputStreams.values()) {
+      Utils.interruptStream(stream);
+    }
     activeInputStreams.clear();
   }
 

src/main/java/me/desair/tus/server/upload/s3/S3StorageService.java

@@ -308,26 +308,38 @@ public UploadInfo append(UploadInfo upload, InputStream inputStream)
     InputStream streamToRead = prepareStreamWithExistingIncompletePart(partObjectKey, inputStream);
 
     // Step 3: Process payload stream in optimal chunk parts and upload to S3
-    AppendResult appendResult =
-        processPayloadChunks(info, streamToRead, info.getId(), partObjectKey);
-
-    // Step 4: Validate minimum append size constraints if configured
-    if (minAppendSize != null && appendResult.totalBytesAppended < minAppendSize) {
-      throw new MinAppendSizeNotMetException(
-          "Append payload size "
-              + appendResult.totalBytesAppended
-              + " is below minimum limit "
-              + minAppendSize);
-    }
+    boolean successfullyFinished = false;
+    try {
+      AppendResult appendResult =
+          processPayloadChunks(info, streamToRead, info.getId(), partObjectKey);
+
+      // Step 4: Validate minimum append size constraints if configured
+      if (minAppendSize != null && appendResult.totalBytesAppended < minAppendSize) {
+        throw new MinAppendSizeNotMetException(
+            "Append payload size "
+                + appendResult.totalBytesAppended
+                + " is below minimum limit "
+                + minAppendSize);
+      }
 
-    // Step 5: Recalculate total uploaded byte offset across all uploaded part objects in S3
-    long newOffset = calculateCurrentOffset(objectKey, info.getId(), partObjectKey);
-    info.setOffset(newOffset);
+      // Step 5: Recalculate total uploaded byte offset across all uploaded part objects in S3
+      long newOffset = calculateCurrentOffset(objectKey, info.getId(), partObjectKey);
+      info.setOffset(newOffset);
+      upload.setOffset(newOffset);
 
-    // Step 6: If all expected bytes are uploaded, compose all part chunks into final S3 object
-    finalizeCompletedUploadIfFinished(info, objectKey, info.getId(), appendResult, newOffset);
-    update(info);
-    return info;
+      // Step 6: If all expected bytes are uploaded, compose all part chunks into final S3 object
+      finalizeCompletedUploadIfFinished(info, objectKey, info.getId(), appendResult, newOffset);
+      update(info);
+      successfullyFinished = true;
+      return info;
+    } finally {
+      if (!successfullyFinished) {
+        long newOffset = calculateCurrentOffset(objectKey, info.getId(), partObjectKey);
+        info.setOffset(newOffset);
+        upload.setOffset(newOffset);
+        update(info);
+      }
+    }
   }
 
   @Override
@@ -678,6 +690,7 @@ private AppendResult processPayloadChunks(
       tempChunkFile.deleteOnExit();
 
       long chunkBytesWritten = 0;
+      IOException readException = null;
       try (FileOutputStream fos = new FileOutputStream(tempChunkFile)) {
         int bytesRead;
         while (chunkBytesWritten < optimalPartSize
@@ -698,13 +711,19 @@ private AppendResult processPayloadChunks(
         if (chunkBytesWritten < optimalPartSize) {
           streamFinished = true;
         }
+      } catch (IOException e) {
+        readException = e;
+        streamFinished = true;
       }
 
       if (chunkBytesWritten == 0) {
         boolean deleted = tempChunkFile.delete();
         if (!deleted) {
           log.warn("Failed to delete temp chunk file {}", tempChunkFile.getAbsolutePath());
         }
+        if (readException != null) {
+          throw readException;
+        }
         break;
       }
 
@@ -721,6 +740,10 @@ private AppendResult processPayloadChunks(
         // Store sub-5MB tail chunk as temporary .part object in S3 for subsequent appends
         storeIncompletePartToS3(partObjectKey, tempChunkFile, chunkBytesWritten);
       }
+
+      if (readException != null) {
+        throw readException;
+      }
     }
 
     return new AppendResult(totalBytesAppended, allPartKeys);

src/test/java/me/desair/tus/server/upload/azure/AzureBlobLockingServiceTest.java

@@ -105,7 +105,9 @@ public void closeInterruptsActiveWatchdogThread() throws Exception {
     InterruptibleInputStream stream = new InterruptibleInputStream(bais);
 
     lockingService.registerInputStream("/test/upload/88888", stream);
-    // KISS: verifying close cleanly shuts down active watchdog thread without exception
+    org.junit.Assert.assertFalse(stream.isInterrupted());
+
     lockingService.close();
+    org.junit.Assert.assertTrue(stream.isInterrupted());
   }
 }