EvoBGP/internal/store/snapshot_budget.go

package store

import (
	"context"
	"os"
	"strings"
	"time"
)

const defaultSnapshotTimeout = 2 * time.Minute

// SnapshotWriteBudget keeps a caller deadline. When ctx has none, it applies
// EVOBGP_SNAPSHOT_TIMEOUT (default 2 minutes). installed is true when this
// function created the deadline; callers use that to set statement_timeout.
func SnapshotWriteBudget(ctx context.Context) (context.Context, context.CancelFunc, bool) {
	if ctx == nil {
		ctx = context.Background()
	}
	if _, ok := ctx.Deadline(); ok {
		return ctx, func() {}, false
	}
	ctx, cancel := context.WithTimeout(ctx, snapshotTimeout())
	return ctx, cancel, true
}

// SnapshotStatementTimeout is the PostgreSQL statement_timeout that matches
// SnapshotWriteBudget. A caller deadline wins; otherwise the env default applies.
func SnapshotStatementTimeout(ctx context.Context) time.Duration {
	if ctx != nil {
		if dl, ok := ctx.Deadline(); ok {
			d := time.Until(dl)
			if d < time.Millisecond {
				return time.Millisecond
			}
			return d
		}
	}
	return snapshotTimeout()
}

func snapshotTimeout() time.Duration {
	s := strings.TrimSpace(os.Getenv("EVOBGP_SNAPSHOT_TIMEOUT"))
	if s == "" {
		return defaultSnapshotTimeout
	}
	d, err := time.ParseDuration(s)
	if err != nil || d <= 0 {
		return defaultSnapshotTimeout
	}
	return d
}