pomerium/internal/zero/reconciler/restart.go

package reconciler

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

// RunWithRestart executes execFn.
// The execution would be restarted, by means of canceling the context provided to execFn, each time restartFn quits.
// the error returned by restartFn is purely informational and does not affect the execution; may be nil.
// the loop is stopped when the context provided to RunWithRestart is canceled or a genuine error is returned by execFn, not caused by the context.
func RunWithRestart(
	ctx context.Context,
	execFn func(context.Context) error,
	restartFn func(context.Context) error,
) error {
	contexts := make(chan context.Context)

	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	var wg sync.WaitGroup
	wg.Add(2)

	var err error
	go func() {
		err = restartWithContext(contexts, execFn)
		cancel()
		wg.Done()
	}()
	go func() {
		restartContexts(ctx, contexts, restartFn)
		wg.Done()
	}()

	wg.Wait()
	return err
}

func restartContexts(
	base context.Context,
	contexts chan<- context.Context,
	restartFn func(context.Context) error,
) {
	defer close(contexts)
	for base.Err() == nil {
		ctx, cancel := context.WithCancelCause(base)
		select {
		case contexts <- ctx:
			err := restartFn(ctx)
			cancel(fmt.Errorf("requesting restart: %w", err))
		case <-base.Done():
			cancel(fmt.Errorf("parent context canceled: %w", base.Err()))
			return
		}
	}
}

func restartWithContext(
	contexts <-chan context.Context,
	execFn func(context.Context) error,
) error {
	var err error
	for ctx := range contexts {
		err = execFn(ctx)
		if ctx.Err() == nil {
			return err
		}
	}
	return err
}