add a proc runner pool

Author: Stumble

pkg/procrunner/pool.go

@@ -0,0 +1,55 @@
+package procrunner
+
+import (
+	"fmt"
+	"sync"
+)
+
+var (
+	ErrMaxReached = fmt.Errorf("max reached")
+)
+
+// ProcRunnerPool is a manager that manages a pool of ProcRunner.
+// It can safely enforce the global memory limit by limiting the number of concurrent ProcRunners.
+type ProcRunnerPool struct {
+	running int
+	max     int
+	mu      sync.Mutex
+}
+
+func NewProcRunnerPool(maxConcurrent int) *ProcRunnerPool {
+	return &ProcRunnerPool{
+		running: 0,
+		max:     maxConcurrent,
+		mu:      sync.Mutex{},
+	}
+}
+
+func (p *ProcRunnerPool) Running() int {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	return p.running
+}
+
+func (p *ProcRunnerPool) Max() int {
+	return p.max
+}
+
+func (p *ProcRunnerPool) NewRunner(filename string, maxheapsizemb uint) (*ProcRunner, error) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	if p.running >= p.max {
+		return nil, ErrMaxReached
+	}
+	runner, err := NewProcRunner(filename, maxheapsizemb)
+	if err != nil {
+		return nil, err
+	}
+	runner.AddPostCloseFn(func() {
+		p.mu.Lock()
+		defer p.mu.Unlock()
+		p.running--
+	})
+	p.running++
+	return runner, nil
+}

pkg/procrunner/pool_test.go

@@ -0,0 +1,41 @@
+package procrunner
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/suite"
+)
+
+type ProcRunnerPoolTestSuite struct {
+	suite.Suite
+}
+
+func TestProcRunnerPoolTestSuite(t *testing.T) {
+	suite.Run(t, new(ProcRunnerPoolTestSuite))
+}
+
+func (suite *ProcRunnerPoolTestSuite) SetupTest() {
+}
+
+func (suite *ProcRunnerPoolTestSuite) TestConcurrency() {
+	pool := NewProcRunnerPool(2)
+	suite.Equal(0, pool.Running())
+	suite.Equal(2, pool.Max())
+	runner1, err := pool.NewRunner("test.js", 32)
+	suite.NoError(err)
+	suite.Equal(1, pool.Running())
+	suite.Equal(2, pool.Max())
+	runner2, err := pool.NewRunner("test.js", 32)
+	suite.NoError(err)
+	suite.Equal(2, pool.Running())
+	runner3, err := pool.NewRunner("test.js", 32)
+	suite.Equal(err, ErrMaxReached)
+	suite.Nil(runner3)
+	runner1.Close()
+	suite.Equal(1, pool.Running())
+	runner4, err := pool.NewRunner("test.js", 32)
+	suite.NoError(err)
+	runner2.Close()
+	runner4.Close()
+	suite.Equal(0, pool.Running())
+}

pkg/procrunner/procrunner.go

@@ -39,6 +39,8 @@ type ProcRunner struct {
 	wg      sync.WaitGroup
 	closeFn func()
 	closed  atomic.Bool
+
+	postCloseFn []func()
 }
 
 // NewProcRunner creates a new ProcRunner that runs the given file.
@@ -87,6 +89,10 @@ func NewProcRunner(fileName string, maxHeapSizeMB uint) (*ProcRunner, error) {
 		defer proc.wg.Done()
 		_ = cmd.Wait()
 		proc.closed.Store(true)
+		// call postCloseFn only after the process is killed
+		for _, f := range proc.postCloseFn {
+			f()
+		}
 	}()
 	// // handle stderr
 	// proc.wg.Add(1)
@@ -191,3 +197,9 @@ func (r *ProcRunner) RunCodeJSON(ctx context.Context, code string) (string, erro
 		return "", ErrorTimeout
 	}
 }
+
+// AddPostCloseFn adds a function to be called after the runner is closed.
+// NOTE: NOT concurrency safe.
+func (r *ProcRunner) AddPostCloseFn(f func()) {
+	r.postCloseFn = append(r.postCloseFn, f)
+}