Add support for multipart uploads in Azure FS

core/trino-main/src/main/java/io/trino/operator/OperatorContext.java

@@ -697,6 +697,18 @@ public ListenableFuture<Void> setBytes(long bytes)
             return blocked;
         }
 
+        @Override
+        public ListenableFuture<Void> addBytes(long delta)
+        {
+            if (delta == 0) {
+                return NOT_BLOCKED;
+            }
+            ListenableFuture<Void> blocked = delegate.addBytes(delta);
+            updateMemoryFuture(blocked, memoryFuture);
+            allocationListener.run();
+            return blocked;
+        }
+
         @Override
         public boolean trySetBytes(long bytes)
         {

core/trino-main/src/test/java/io/trino/execution/buffer/TestSpoolingExchangeOutputBuffer.java

@@ -415,6 +415,12 @@ public ListenableFuture<Void> setBytes(long bytes)
             return immediateVoidFuture();
         }
 
+        @Override
+        public ListenableFuture<Void> addBytes(long delta)
+        {
+            return immediateVoidFuture();
+        }
+
         @Override
         public boolean trySetBytes(long bytes)
         {

lib/trino-memory-context/src/main/java/io/trino/memory/context/CoarseGrainLocalMemoryContext.java

@@ -20,6 +20,7 @@
 import com.google.errorprone.annotations.concurrent.GuardedBy;
 
 import static com.google.common.base.Preconditions.checkArgument;
+import static java.lang.Math.addExact;
 import static java.util.Objects.requireNonNull;
 
 /**
@@ -69,6 +70,12 @@ public synchronized ListenableFuture<Void> setBytes(long bytes)
         return Futures.immediateVoidFuture();
     }
 
+    @Override
+    public synchronized ListenableFuture<Void> addBytes(long delta)
+    {
+        return setBytes(addExact(currentBytes, delta));
+    }
+
     @Override
     public synchronized boolean trySetBytes(long bytes)
     {

lib/trino-memory-context/src/main/java/io/trino/memory/context/LocalMemoryContext.java

@@ -30,6 +30,17 @@ public interface LocalMemoryContext
      */
     ListenableFuture<Void> setBytes(long bytes);
 
+    /**
+     * When this method returns, the bytes tracked by this LocalMemoryContext has been updated by the provided delta.
+     * The returned future will tell the caller whether it should block before reserving more memory
+     * (which happens when the memory pools are low on memory).
+     * <p>
+     * Note: Canceling the returned future will complete it immediately even though the memory pools are low
+     * on memory, and callers blocked on this future will proceed to allocating more memory from the exhausted
+     * pools, which will violate the protocol of Trino MemoryPool implementation.
+     */
+    ListenableFuture<Void> addBytes(long delta);
+
     /**
      * This method can return false when there is not enough memory available to satisfy a positive delta allocation
      * ({@code bytes} is greater than the bytes tracked by this LocalMemoryContext).

lib/trino-memory-context/src/main/java/io/trino/memory/context/SimpleLocalMemoryContext.java

@@ -67,6 +67,12 @@ public synchronized ListenableFuture<Void> setBytes(long bytes)
         return future;
     }
 
+    @Override
+    public synchronized ListenableFuture<Void> addBytes(long delta)
+    {
+        return setBytes(usedBytes + delta);
+    }
+
     @Override
     public synchronized boolean trySetBytes(long bytes)
     {