EvoBGP/internal/jobs/context.go

package jobs

import (
	"context"
	"os"
	"strconv"
	"time"
)

const (
	defaultJobTimeoutModuleRefresh = 10 * time.Minute
	defaultJobTimeoutTenantRefresh = 15 * time.Minute
	defaultJobTimeoutDeployApply   = 5 * time.Minute
	defaultJobTimeoutPeerReconcile = 10 * time.Minute
	defaultJobTimeoutRollback      = 5 * time.Minute
	defaultJobTimeoutBirdReload    = 2 * time.Minute
)

func jobTimeout(kind string) time.Duration {
	envKey := map[string]string{
		KindModuleRefresh:    "EVOBGP_JOB_TIMEOUT_MODULE_REFRESH",
		KindTenantRefresh:    "EVOBGP_JOB_TIMEOUT_TENANT_REFRESH",
		KindDeployApply:      "EVOBGP_JOB_TIMEOUT_DEPLOY_APPLY",
		KindPeerReconcile:    "EVOBGP_JOB_TIMEOUT_PEER_RECONCILE",
		KindRevisionRollback: "EVOBGP_JOB_TIMEOUT_ROLLBACK",
		KindBirdReload:       "EVOBGP_JOB_TIMEOUT_BIRD_RELOAD",
	}[kind]
	if envKey != "" {
		if d, err := time.ParseDuration(os.Getenv(envKey)); err == nil && d > 0 {
			return d
		}
	}
	switch kind {
	case KindModuleRefresh:
		return defaultJobTimeoutModuleRefresh
	case KindTenantRefresh:
		return defaultJobTimeoutTenantRefresh
	case KindDeployApply:
		return defaultJobTimeoutDeployApply
	case KindPeerReconcile:
		return defaultJobTimeoutPeerReconcile
	case KindRevisionRollback:
		return defaultJobTimeoutRollback
	case KindBirdReload:
		return defaultJobTimeoutBirdReload
	default:
		if n, err := strconv.Atoi(os.Getenv("EVOBGP_JOB_TIMEOUT_SEC")); err == nil && n > 0 {
			return time.Duration(n) * time.Second
		}
		return defaultJobTimeoutModuleRefresh
	}
}

// workContext returns a timeout context that also cancels when the job is cancelled.
func (j *Job) workContext() (context.Context, context.CancelFunc) {
	if j == nil {
		return context.Background(), func() {}
	}
	timeout := jobTimeout(j.Kind)
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	go func() {
		ticker := time.NewTicker(500 * time.Millisecond)
		defer ticker.Stop()
		for {
			select {
			case <-ctx.Done():
				return
			case <-ticker.C:
				if j.IsCancelRequested() {
					cancel()
					return
				}
			}
		}
	}()
	return ctx, cancel
}