depot/go/nix/nixpool/nixpool.go

// SPDX-FileCopyrightText: 2023 Luke Granger-Brown <depot@lukegb.com>
//
// SPDX-License-Identifier: Apache-2.0

package nixpool

import (
	"context"
	"log"
	"math/rand"
	"runtime"
	"sync"
	"time"

	"git.lukegb.com/lukegb/depot/go/nix/nixstore"
)

type Pool struct {
	factory DaemonFactory

	cond      *sync.Cond
	limit     int
	generated int
	available []*nixstore.Daemon
}

func New(ctx context.Context, factory DaemonFactory, limit int) *Pool {
	return &Pool{
		factory: func() (*nixstore.Daemon, error) {
			var d *nixstore.Daemon
			var err error
			retryInterval := 10 * time.Millisecond
			maxRetry := 10 * time.Second
			for n := 0; n < 20; n++ {
				d, err = factory()
				if err == nil {
					break
				}
				log.Printf("failed to connect: %v", err)
				time.Sleep(retryInterval + time.Duration((rand.Float64()-0.5)*float64(retryInterval)))
				retryInterval = retryInterval * 2
				if retryInterval > maxRetry {
					retryInterval = maxRetry
				}
			}
			return d, err
		},

		cond:  sync.NewCond(new(sync.Mutex)),
		limit: limit,
	}
}

func (p *Pool) getFromAvailableLocked() *nixstore.Daemon {
	end := len(p.available) - 1
	d := p.available[end]
	p.available = p.available[:end]
	return d
}

func (p *Pool) log(ln string) {
	if false {
		_, file, line, _ := runtime.Caller(2)
		log.Printf("%p (%s:%d): %s", p, file, line, ln)
	}
}

func (p *Pool) Get() (*nixstore.Daemon, error) {
	p.cond.L.Lock()

	for {
		if len(p.available) > 0 {
			d := p.getFromAvailableLocked()
			p.log("pool.Get: from available")
			p.cond.L.Unlock()
			return d, nil
		}

		if p.generated < p.limit {
			p.generated++
			p.cond.L.Unlock()
			d, err := p.factory()
			if err != nil {
				p.cond.L.Lock()
				p.generated--
				p.cond.L.Unlock()
				return nil, err
			}
			p.log("pool.Get: generated new")
			return d, nil
		}

		p.cond.Wait()
	}
}

func (p *Pool) Put(d *nixstore.Daemon) {
	p.cond.L.Lock()
	if d.Err() != nil {
		d.Close()
		p.log("pool.Put: broken")
		p.generated--
	} else {
		p.log("pool.Put: returned to pool")
		p.available = append(p.available, d)
	}
	p.cond.Signal()
	p.cond.L.Unlock()
}

func (p *Pool) Busyness() float64 {
	p.cond.L.Lock()
	defer p.cond.L.Unlock()

	currentlyAvailable := len(p.available)
	currentlyCheckedOut := p.generated - currentlyAvailable
	proportionCheckedOut := float64(currentlyCheckedOut) / float64(p.limit)
	return proportionCheckedOut
}