Feature Request: Daemon/Sidekick Pool

[sbstp]

I have this pattern where I have a main goroutine which is accompanied/supported by side goroutines. The main goroutines produces a value, whereas the side goroutines do not produce any values to capture but perform some sort of upkeep or utility function. Once the main goroutine exits, the side goroutines' context should be canceled. If the side goroutines exit with an error or panic, the main and other side goroutines should be canceled.

An example of where I use this pattern is inside of gRPC server-streaming handlers. I send a keep alive message on a regular interval to keep the connection alive while the main goroutine is waiting on an event to happen. Once the main goroutine is done I want to kill the keep-alive routine obviously.

Naming is hard, I'm calling this the daemon/sidekick pattern for now, similar to daemon threads in Python. It's also kind of similar to sidecar pods in Kubernetes.

I have a basic implementation of this here https://github.com/sbstp/unison/blob/master/sidekick.go but it's a bit primitive and doesn't print panics cleanly.

Right now Pool.Wait waits for all the routines to be done, if there was a way to wait only for the first routine to be finished and cancel the others, this pattern could be achieved easily.

Is there interest in adding feature support for this?

[ljluestc]

package conc


import (
	"context"
	"fmt"
	"runtime"
	"sync"
)

// SidekickPool manages a main goroutine and side goroutines with the daemon/sidekick pattern.
type SidekickPool[T any] struct {
	ctx        context.Context
	cancel     func()
	wg         sync.WaitGroup
	mainTask   chan func() (T, error)
	result     chan result[T]
	mainDone   bool
	mu         sync.Mutex
	err        error
	panicC     PanicCatcher
}

// result holds the main goroutine’s output and error.
type result[T any] struct {
	value T
	err   error
}

// NewSidekickPool creates a new SidekickPool for the daemon/sidekick pattern.
func NewSidekickPool[T any]() *SidekickPool[T] {
	ctx, cancel := context.WithCancel(context.Background())
	return &SidekickPool[T]{
		ctx:      ctx,
		cancel:   cancel,
		mainTask: make(chan func() (T, error), 1),
		result:   make(chan result[T], 1),
	}
}

// GoMain submits the main goroutine, which produces a value of type T.
func (p *SidekickPool[T]) GoMain(f func() (T, error)) {
	p.mu.Lock()
	if p.mainDone {
		p.mu.Unlock()
		panic("GoMain called after main task already submitted")
	}
	p.mainDone = true
	p.mu.Unlock()

	p.wg.Add(1)
	p.mainTask <- f
	go func() {
		defer p.wg.Done()
		defer p.cancel() // Cancel side goroutines when main exits.

		f := <-p.mainTask
		p.panicC.Try(func() {
			value, err := f()
			p.result <- result[T]{value, err}
		})
		if err := p.panicC.Err(); err != nil {
			p.mu.Lock()
			p.err = fmt.Errorf("main goroutine panicked: %w", err)
			p.mu.Unlock()
		}
	}()
}

// Go submits a side goroutine, which runs until canceled or an error occurs.
func (p *SidekickPool[T]) Go(f func() error) {
	p.wg.Add(1)
	go func() {
		defer p.wg.Done()
		defer p.cancel() // Cancel all goroutines if side task errors or panics.

		p.panicC.Try(func() {
			if err := f(); err != nil {
				p.mu.Lock()
				p.err = fmt.Errorf("side goroutine failed: %w", err)
				p.mu.Unlock()
			}
		})
		if err := p.panicC.Err(); err != nil {
			p.mu.Lock()
			p.err = fmt.Errorf("side goroutine panicked: %w", err)
			p.mu.Unlock()
		}
	}()
}

// Wait waits for the main goroutine to complete, cancels side goroutines, and returns the main result.
func (p *SidekickPool[T]) Wait() (T, error) {
	// Wait for the main goroutine to produce a result.
	select {
	case res := <-p.result:
		// Cancel side goroutines.
		p.cancel()
		// Wait for all goroutines to finish.
		p.wg.Wait()
		// Check for errors from side goroutines or panics.
		p.mu.Lock()
		if p.err != nil {
			p.mu.Unlock()
			return res.value, p.err
		}
		p.mu.Unlock()
		return res.value, res.err
	case <-p.ctx.Done():
		// If context is canceled (e.g., by a side goroutine), wait for cleanup.
		p.wg.Wait()
		p.mu.Lock()
		err := p.err
		if err == nil {
			err = context.Canceled
		}
		p.mu.Unlock()
		var zero T
		return zero, err
	}
}

// PanicCatcher captures panics and converts them to errors with stack traces.
type PanicCatcher struct {
	err error
}

// Try executes f and captures any panic.
func (pc *PanicCatcher) Try(f func()) {
	defer func() {
		if r := recover(); r != nil {
			const size = 64 << 10
			buf := make([]byte, size)
			buf = buf[:runtime.Stack(buf, false)]
			pc.err = fmt.Errorf("panic: %v\n%s", r, buf)
		}
	}()
	f()
}

// Err returns the captured panic error, if any.
func (pc *PanicCatcher) Err() error {
	return pc.err
}