diff --git a/.github/workflows/go.yml b/.github/workflows/go.yml index c307c50..b673694 100644 --- a/.github/workflows/go.yml +++ b/.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 diff --git a/permitpool/example_test.go b/permitpool/example_test.go new file mode 100644 index 0000000..9c5a860 --- /dev/null +++ b/permitpool/example_test.go @@ -0,0 +1,44 @@ +package permitpool_test + +import ( + "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) + + 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() + // 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 +} diff --git a/permitpool/go.mod b/permitpool/go.mod new file mode 100644 index 0000000..12c7d7c --- /dev/null +++ b/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 +) diff --git a/permitpool/go.sum b/permitpool/go.sum new file mode 100644 index 0000000..713a0b4 --- /dev/null +++ b/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= diff --git a/permitpool/permitpool.go b/permitpool/permitpool.go new file mode 100644 index 0000000..be38d60 --- /dev/null +++ b/permitpool/permitpool.go @@ -0,0 +1,41 @@ +// Package permitpool exposes a synchronization primitive for +// limiting the number of concurrent operations. See the Pool +// example for a simple use case. +package permitpool + +// 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 +func (c *Pool) Acquire() { + c.sem <- struct{}{} +} + +// 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) +} diff --git a/permitpool/permitpool_test.go b/permitpool/permitpool_test.go new file mode 100644 index 0000000..fc1165a --- /dev/null +++ b/permitpool/permitpool_test.go @@ -0,0 +1,59 @@ +package permitpool_test + +import ( + "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() + pool := permitpool.New(2) + require.NotNil(t, pool) + assert.Equal(t, 0, pool.CurrentPermits(), "Expected 0 permits initially") + + pool.Acquire() + assert.Equal(t, 1, pool.CurrentPermits(), "Expected 1 permit after Acquire") + + pool.Acquire() + 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() + pool.Acquire() + + start := make(chan struct{}) + testChan := make(chan struct{}) + go func() { + close(start) + pool.Acquire() + 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") + } +}