permitpool: add new module

.github/workflows/go.yml

@@ -8,27 +8,29 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        module: ["awsutil",
-                  "base62",
-                  "configutil",
-                  "cryptoutil",
-                  "fileutil",
-                  "gatedwriter",
-                  "httputil",
-                  "kv-builder",
-                  "listenerutil",
-                  "mlock",
-                  "nonceutil",
-                  "parseutil",
-                  "password",
-                  "plugincontainer",
-                  "pluginutil",
-                  "random",
-                  "reloadutil",
-                  "strutil",
-                  "temperror",
-                  "tlsutil",
-                  "toggledlogger"]
+        module:
+          - "awsutil"
+          - "base62"
+          - "configutil"
+          - "cryptoutil"
+          - "fileutil"
+          - "gatedwriter"
+          - "httputil"
+          - "kv-builder"
+          - "listenerutil"
+          - "mlock"
+          - "nonceutil"
+          - "parseutil"
+          - "password"
+          - "permitpool"
+          - "plugincontainer"
+          - "pluginutil"
+          - "random"
+          - "reloadutil"
+          - "strutil"
+          - "temperror"
+          - "tlsutil"
+          - "toggledlogger"
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@c85c95e3d7251135ab7dc9ce3241c5835cc595a9 # v3.5.3

permitpool/example_test.go

@@ -0,0 +1,46 @@
+package permitpool_test
+
+import (
+	"context"
+	"crypto/ecdsa"
+	"crypto/elliptic"
+	"crypto/rand"
+	"fmt"
+	"sync"
+
+	"github.com/hashicorp/go-secure-stdlib/permitpool"
+)
+
+func ExamplePool() {
+	// Create a new permit pool with 2 permits.
+	// This limits the number of concurrent operations to 2.
+	pool := permitpool.New(2)
+	ctx := context.Background()
+
+	keys := make([]*ecdsa.PrivateKey, 5)
+	wg := &sync.WaitGroup{}
+	for i := range 5 {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			// Acquire a permit from the pool. This
+			// will block until a permit is available
+			// and assigned to this goroutine.
+			pool.Acquire(ctx)
+			// Ensure the permit is returned to the pool upon
+			// completion of the operation.
+			defer pool.Release()
+
+			// Perform some expensive operation
+			key, err := ecdsa.GenerateKey(elliptic.P384(), rand.Reader)
+			if err != nil {
+				fmt.Println("Failed to generate key:", err)
+				return
+			}
+			keys[i] = key
+		}()
+	}
+	wg.Wait()
+	fmt.Printf("Generated %d keys\n", len(keys))
+	// Output: Generated 5 keys
+}

permitpool/go.mod

@@ -0,0 +1,11 @@
+module github.com/hashicorp/go-secure-stdlib/permitpool
+
+go 1.23
+
+require github.com/stretchr/testify v1.10.0
+
+require (
+	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
+)

permitpool/go.sum

@@ -0,0 +1,10 @@
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA=
+github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

permitpool/permitpool.go

@@ -0,0 +1,49 @@
+// Package permitpool exposes a synchronization primitive for
+// limiting the number of concurrent operations. See the Pool
+// example for a simple use case.
+package permitpool
+
+import "context"
+
+// DefaultParallelOperations is the default number of parallel operations
+// allowed by the permit pool.
+const DefaultParallelOperations = 128
+
+// Pool is used to limit maximum outstanding requests
+type Pool struct {
+	sem chan struct{}
+}
+
+// New returns a new permit pool with the provided
+// number of permits. If permits is less than 1, the
+// default number of parallel operations is used.
+func New(permits int) *Pool {
+	if permits < 1 {
+		permits = DefaultParallelOperations
+	}
+	return &Pool{
+		sem: make(chan struct{}, permits),
+	}
+}
+
+// Acquire returns when a permit has been acquired, or
+// if the context is canceled.
+func (c *Pool) Acquire(ctx context.Context) error {
+	select {
+	case c.sem <- struct{}{}:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}
+
+// Release returns a permit to the pool
+func (c *Pool) Release() {
+	<-c.sem
+}
+
+// CurrentPermits gets the number of used permits.
+// This corresponds to the number of running operations.
+func (c *Pool) CurrentPermits() int {
+	return len(c.sem)
+}

permitpool/permitpool_test.go

@@ -0,0 +1,112 @@
+package permitpool_test
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/hashicorp/go-secure-stdlib/permitpool"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestPermitPool(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	pool := permitpool.New(2)
+	require.NotNil(t, pool)
+	assert.Equal(t, 0, pool.CurrentPermits(), "Expected 0 permits initially")
+
+	pool.Acquire(ctx)
+	assert.Equal(t, 1, pool.CurrentPermits(), "Expected 1 permit after Acquire")
+
+	pool.Acquire(ctx)
+	assert.Equal(t, 2, pool.CurrentPermits(), "Expected 2 permits after second Acquire")
+
+	pool.Release()
+	assert.Equal(t, 1, pool.CurrentPermits(), "Expected 1 permit after Release")
+
+	pool.Release()
+	assert.Equal(t, 0, pool.CurrentPermits(), "Expected 0 permits after second Release")
+
+	pool.Acquire(ctx)
+	pool.Acquire(ctx)
+
+	start := make(chan struct{})
+	testChan := make(chan struct{})
+	go func() {
+		close(start)
+		pool.Acquire(ctx)
+		defer pool.Release()
+		close(testChan)
+	}()
+
+	// Wait for the goroutine to start
+	<-start
+	select {
+	case <-testChan:
+		t.Error("Expected Acquire when no permits available to block")
+	case <-time.After(10 * time.Millisecond):
+		// Success, the goroutine is blocked
+	}
+
+	pool.Release()
+	pool.Release()
+	select {
+	case <-testChan:
+		// Success, the goroutine has acquired the permit
+	case <-time.After(10 * time.Millisecond):
+		t.Error("Expected Acquire to unblock when a permit is available")
+	}
+}
+
+func TestAcquireContextCancellation(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	pool := permitpool.New(2)
+	require.NotNil(t, pool)
+
+	// Acquire all permits
+	pool.Acquire(ctx)
+	pool.Acquire(ctx)
+
+	// Test AcquireContext blocks until context is canceled or a permit is available
+	testChan := make(chan struct{})
+	ctx, cancel := context.WithCancel(context.Background())
+	go func() {
+		defer close(testChan)
+		err := pool.Acquire(ctx)
+		if err != nil {
+			assert.ErrorIs(t, err, context.Canceled)
+			return
+		}
+		pool.Release()
+	}()
+
+	select {
+	case <-testChan:
+		t.Error("Expected AcquireContext to block until context is canceled")
+	case <-time.After(10 * time.Millisecond):
+		// Success, the goroutine is blocked
+	}
+
+	cancel()
+
+	select {
+	case <-testChan:
+		// Success, the goroutine errored out wit a canceled context
+	case <-time.After(10 * time.Millisecond):
+		t.Error("Expected AcquireContext to unblock when context is canceled")
+	}
+
+	// Make one permit available
+	pool.Release()
+
+	err := pool.Acquire(context.Background())
+	require.NoError(t, err)
+
+	pool.Release()
+	pool.Release()
+}