From 9a84a0964c85f3bd8abbd921c9a4a86886acca34 Mon Sep 17 00:00:00 2001 From: Johan Brandhorst-Satzkorn Date: Thu, 9 Jan 2025 09:38:50 -0800 Subject: [PATCH] permitpool: add new module The permitpool module contains a single package that implements a pool of permits, allowing limited number of concurrent executions. --- .github/workflows/go.yml | 44 +++++++++++++------------- permitpool/go.mod | 11 +++++++ permitpool/go.sum | 10 ++++++ permitpool/permitpool.go | 38 ++++++++++++++++++++++ permitpool/permitpool_test.go | 59 +++++++++++++++++++++++++++++++++++ 5 files changed, 141 insertions(+), 21 deletions(-) create mode 100644 permitpool/go.mod create mode 100644 permitpool/go.sum create mode 100644 permitpool/permitpool.go create mode 100644 permitpool/permitpool_test.go 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/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..acd4fa5 --- /dev/null +++ b/permitpool/permitpool.go @@ -0,0 +1,38 @@ +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") + } +}