pool.go

package nodejs

import (
	"context"
	"runtime"
	"sync"
	"sync/atomic"
	"time"
)

// Pool manages a pool of NodeJS processes that can be used to execute JavaScript code.
// It automatically maintains a queue of ready processes and creates new ones as needed.
type Pool struct {
	queue     []*Process    // Queue of available NodeJS processes
	factory   *Factory      // Factory used to create new processes
	maxProcs  uint64        // Maximum number of processes allowed
	procs     uint64        // Current count of processes
	spawner   atomic.Bool   // Flag indicating if spawner goroutine is running
	addcond   *sync.Cond    // Condition variable for process addition
	availcond *sync.Cond    // Condition variable for process availability
	lk        sync.Mutex    // Mutex for protecting shared state
	Timeout   time.Duration // Timeout for running nodejs, defaults to 10 second if zero
}

// NewPool returns a pool of nodejs processes generated by the factory. The pool will
// always generate at least queueSize processes in memory and will start new ones as
// they are taken. Passing queueSize <=0 will use the number of CPUs.
func (f *Factory) NewPool(queueSize, maxProcs int) *Pool {
	if queueSize <= 0 {
		queueSize = runtime.NumCPU()
	}
	if maxProcs <= 0 {
		maxProcs = runtime.NumCPU()
	}

	p := &Pool{
		maxProcs: uint64(maxProcs),
		factory:  f,
	}
	p.addcond = sync.NewCond(&p.lk)
	p.availcond = sync.NewCond(&p.lk)

	go p.run()

	return p
}

// TakeIfAvailable returns a Process if any is immediately available, or nil if not.
// This method does not wait or block if no processes are available.
func (pool *Pool) TakeIfAvailable() *Process {
	pool.lk.Lock()
	defer pool.lk.Unlock()

	if len(pool.queue) == 0 {
		return nil
	}
	res := pool.queue[0]
	pool.queue = pool.queue[1:]
	return res
}

// Take returns a [Process] from the pool and will wait until one is available,
// unless the context is cancelled.
func (pool *Pool) Take(ctx context.Context) (*Process, error) {
	pool.lk.Lock()
	defer pool.lk.Unlock()

	def := context.AfterFunc(ctx, func() {
		pool.availcond.L.Lock()
		defer pool.availcond.L.Unlock()
		pool.availcond.Broadcast()
	})
	defer def()

	for {
		if e := ctx.Err(); e != nil {
			return nil, e
		}
		if len(pool.queue) > 0 {
			res := pool.queue[0]
			pool.queue = pool.queue[1:]
			return res, nil
		}
		pool.availcond.Wait()
	}
}

// TakeTimeout will return a [Process] taken from the pool if any is available
// before the timeout expires.
func (pool *Pool) TakeTimeout(t time.Duration) (*Process, error) {
	ctx, cancel := context.WithTimeout(context.Background(), t)
	defer cancel()

	return pool.Take(ctx)
}

// addProcess adds a new Process to the pool and notifies waiting goroutines
// that a new process is available.
func (pool *Pool) addProcess(p *Process) {
	pool.lk.Lock()
	defer pool.lk.Unlock()

	pool.queue = append(pool.queue, p)
	pool.availcond.Broadcast() // Signal that a process is available
}

// run is the main goroutine that creates new NodeJS processes for the pool.
// It ensures only one spawner is running at a time and maintains the number
// of processes within limits.
func (pool *Pool) run() {
	// Ensure only one spawner is running
	if !pool.spawner.CompareAndSwap(false, true) {
		// means run() is being called more than once
		return
	}
	defer pool.spawner.Store(false)

	for {
		// Check if we've reached the maximum allowed processes
		if atomic.LoadUint64(&pool.procs) > pool.maxProcs {
			// too many running, stop
			// Note: we can have a very special edge case where we end at the exact time all processes stop
			// this is very unlikely to happen, especially if we have multiple processes
			return
		}

		// Use configured timeout or default to 10 seconds
		timeout := pool.Timeout
		if timeout <= 0 {
			timeout = 10 * time.Second
		}

		// Create a new NodeJS process
		newProc, err := pool.factory.NewWithTimeout(timeout)
		if err != nil {
			// If creation fails, wait and try again
			time.Sleep(10 * time.Second)
			continue
		}

		// Increment process counter and register cleanup handler
		atomic.AddUint64(&pool.procs, 1)
		newProc.cleanup = append(newProc.cleanup, pool.procsSubOne)

		// Add the new process to the available queue
		pool.addProcess(newProc)
	}
}

// procsSubOne decrements the process counter and restarts the spawner if needed.
// This is called when a process exits or is closed, as part of its cleanup handlers.
func (pool *Pool) procsSubOne() {
	// Decrement the process counter (^uint64(0) is the bitwise NOT of 0, which is all 1s)
	// In 2's complement, adding all 1s is equivalent to subtracting 1
	atomic.AddUint64(&pool.procs, ^uint64(0))

	// If the spawner isn't running, restart it to maintain the pool
	if !pool.spawner.Load() {
		go pool.run()
	}
}