Add regression test for semaphore permits after a failed block batch

Drive the real uploadAndCollectBlockData path against a local HTTP test
server that fails one block of the batch, with blockUploadSemaphore sized
to 2, and assert that both permits can be acquired again afterwards.

On master (9d772d0) this fails with "context deadline exceeded" because
the collector returns on the first error and the remaining worker stays
blocked on the unbuffered result channel, never releasing its permit.
With the worker-drain change in this PR it passes, including under -race.

Author: jomplox

file_upload_batch_permits_test.go

@@ -0,0 +1,123 @@
+package proton_api_bridge
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/ProtonMail/gopenpgp/v3/crypto"
+	"github.com/rclone/go-proton-api"
+	"golang.org/x/sync/semaphore"
+)
+
+func TestFailedBlockBatchReleasesEverySemaphorePermit(t *testing.T) {
+	originalBlockSize := UPLOAD_BLOCK_SIZE
+	originalBatchSize := UPLOAD_BATCH_BLOCK_SIZE
+	UPLOAD_BLOCK_SIZE = 16
+	UPLOAD_BATCH_BLOCK_SIZE = 2
+	t.Cleanup(func() {
+		UPLOAD_BLOCK_SIZE = originalBlockSize
+		UPLOAD_BATCH_BLOCK_SIZE = originalBatchSize
+	})
+
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Date", time.Now().UTC().Format(http.TimeFormat))
+		w.Header().Set("Content-Type", "application/json")
+
+		switch r.URL.Path {
+		case "/drive/blocks":
+			_, err := io.WriteString(w, fmt.Sprintf(
+				`{"Code":1000,"UploadLinks":[{"Token":"error-token","BareURL":%q},{"Token":"success-token","BareURL":%q}]}`,
+				serverURL(r, "/storage/blocks/error"),
+				serverURL(r, "/storage/blocks/success"),
+			))
+			if err != nil {
+				t.Error(err)
+			}
+		case "/storage/blocks/error":
+			w.WriteHeader(http.StatusBadGateway)
+			_, err := io.WriteString(w, `{"Code":0,"Error":"simulated bad gateway"}`)
+			if err != nil {
+				t.Error(err)
+			}
+		case "/storage/blocks/success":
+			time.Sleep(100 * time.Millisecond)
+			_, err := io.WriteString(w, `{"Code":1000}`)
+			if err != nil {
+				t.Error(err)
+			}
+		default:
+			http.NotFound(w, r)
+		}
+	}))
+	defer server.Close()
+
+	manager := proton.New(
+		proton.WithHostURL(server.URL),
+		proton.WithRetryCount(0),
+	)
+	defer manager.Close()
+	client := manager.NewClient("", "", "")
+	defer client.Close()
+
+	pgp := crypto.PGP()
+	signingKey, err := pgp.KeyGeneration().AddUserId("test", "test@example.com").New().GenerateKey()
+	if err != nil {
+		t.Fatal(err)
+	}
+	signingKeyRing, err := crypto.NewKeyRing(signingKey)
+	if err != nil {
+		t.Fatal(err)
+	}
+	nodeKey, err := pgp.KeyGeneration().AddUserId("node", "node@example.com").New().GenerateKey()
+	if err != nil {
+		t.Fatal(err)
+	}
+	nodeKeyRing, err := crypto.NewKeyRing(nodeKey)
+	if err != nil {
+		t.Fatal(err)
+	}
+	sessionKey, err := pgp.GenerateSessionKey()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	blockSemaphore := semaphore.NewWeighted(2)
+	drive := &ProtonDrive{
+		MainShare: &proton.Share{
+			ShareMetadata: proton.ShareMetadata{ShareID: "share-id"},
+			AddressID:     "address-id",
+		},
+		DefaultAddrKR:        signingKeyRing,
+		c:                    client,
+		blockUploadSemaphore: blockSemaphore,
+	}
+
+	_, _, _, _, err = drive.uploadAndCollectBlockData(
+		context.Background(),
+		sessionKey,
+		nodeKeyRing,
+		bytes.NewReader([]byte("0123456789abcdef0123456789abcdef")),
+		"link-id",
+		"revision-id",
+	)
+	if err == nil {
+		t.Fatal("expected failed block upload")
+	}
+
+	acquireCtx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+	if err := blockSemaphore.Acquire(acquireCtx, 2); err != nil {
+		t.Fatalf("failed batch leaked a block-upload permit: %v", err)
+	}
+	blockSemaphore.Release(2)
+}
+
+func serverURL(r *http.Request, path string) string {
+	return "http://" + r.Host + path
+}