freedakgmail
139 lines
4.1 KiB
Go
139 lines
4.1 KiB
Go
package connector
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"time"
|
|
|
|
"github.com/edgeai/gateway/internal/config"
|
|
)
|
|
|
|
// TimeoutManager manages layered timeouts for different phases of request processing.
|
|
type TimeoutManager struct {
|
|
cfg *config.TimeoutConfig
|
|
}
|
|
|
|
// NewTimeoutManager creates a new TimeoutManager.
|
|
func NewTimeoutManager(cfg *config.TimeoutConfig) *TimeoutManager {
|
|
return &TimeoutManager{cfg: cfg}
|
|
}
|
|
|
|
// TimeoutPhase represents a phase of request processing.
|
|
type TimeoutPhase string
|
|
|
|
const (
|
|
PhaseQueue TimeoutPhase = "queue"
|
|
PhaseFirstToken TimeoutPhase = "first_token"
|
|
PhaseInference TimeoutPhase = "inference"
|
|
PhaseTotal TimeoutPhase = "total"
|
|
)
|
|
|
|
// TimeoutConfig holds resolved timeout values for a specific request.
|
|
type TimeoutConfig struct {
|
|
QueueMs int
|
|
FirstTokenMs int
|
|
InferenceMs int
|
|
TotalMs int
|
|
ConnectMs int
|
|
IdleMs int
|
|
}
|
|
|
|
// ResolveTimeouts merges request-level timeout overrides with global defaults.
|
|
func (tm *TimeoutManager) ResolveTimeouts(reqTimeouts *config.TimeoutConfig, overrides map[string]int) *TimeoutConfig {
|
|
tc := &TimeoutConfig{
|
|
QueueMs: tm.cfg.DefaultQueueMs,
|
|
FirstTokenMs: tm.cfg.DefaultFirstTokenMs,
|
|
InferenceMs: tm.cfg.DefaultInferenceMs,
|
|
TotalMs: tm.cfg.DefaultTotalMs,
|
|
ConnectMs: tm.cfg.DefaultConnectMs,
|
|
IdleMs: tm.cfg.DefaultIdleMs,
|
|
}
|
|
|
|
if overrides != nil {
|
|
if v, ok := overrides["queue_ms"]; ok && v > 0 {
|
|
tc.QueueMs = v
|
|
}
|
|
if v, ok := overrides["first_token_ms"]; ok && v > 0 {
|
|
tc.FirstTokenMs = v
|
|
}
|
|
if v, ok := overrides["inference_ms"]; ok && v > 0 {
|
|
tc.InferenceMs = v
|
|
}
|
|
if v, ok := overrides["total_ms"]; ok && v > 0 {
|
|
tc.TotalMs = v
|
|
}
|
|
}
|
|
|
|
return tc
|
|
}
|
|
|
|
// QueueContext returns a context with the queue timeout.
|
|
func (tm *TimeoutManager) QueueContext(parent context.Context, tc *TimeoutConfig) (context.Context, context.CancelFunc) {
|
|
return context.WithTimeout(parent, time.Duration(tc.QueueMs)*time.Millisecond)
|
|
}
|
|
|
|
// InferenceContext returns a context with the inference timeout.
|
|
func (tm *TimeoutManager) InferenceContext(parent context.Context, tc *TimeoutConfig) (context.Context, context.CancelFunc) {
|
|
return context.WithTimeout(parent, time.Duration(tc.InferenceMs)*time.Millisecond)
|
|
}
|
|
|
|
// TotalContext returns a context with the total request timeout.
|
|
func (tm *TimeoutManager) TotalContext(parent context.Context, tc *TimeoutConfig) (context.Context, context.CancelFunc) {
|
|
return context.WithTimeout(parent, time.Duration(tc.TotalMs)*time.Millisecond)
|
|
}
|
|
|
|
// CheckTimeout returns an error if the given phase has timed out.
|
|
func (tm *TimeoutManager) CheckTimeout(phase TimeoutPhase, elapsed time.Duration, tc *TimeoutConfig) error {
|
|
var limit time.Duration
|
|
switch phase {
|
|
case PhaseQueue:
|
|
limit = time.Duration(tc.QueueMs) * time.Millisecond
|
|
case PhaseFirstToken:
|
|
limit = time.Duration(tc.FirstTokenMs) * time.Millisecond
|
|
case PhaseInference:
|
|
limit = time.Duration(tc.InferenceMs) * time.Millisecond
|
|
case PhaseTotal:
|
|
limit = time.Duration(tc.TotalMs) * time.Millisecond
|
|
default:
|
|
return nil
|
|
}
|
|
|
|
if elapsed > limit {
|
|
return fmt.Errorf("%s timeout: elapsed %v exceeds limit %v", phase, elapsed, limit)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// CancelManager manages cancellation propagation from client to inference engine.
|
|
type CancelManager struct{}
|
|
|
|
// NewCancelManager creates a new CancelManager.
|
|
func NewCancelManager() *CancelManager {
|
|
return &CancelManager{}
|
|
}
|
|
|
|
// WatchClientDisconnect watches for client connection close and signals cancellation.
|
|
// Returns a context that is cancelled when the client disconnects.
|
|
func (cm *CancelManager) WatchClientDisconnect(r interface{ Done() <-chan struct{} }, cancel context.CancelFunc) {
|
|
go func() {
|
|
select {
|
|
case <-r.Done():
|
|
cancel()
|
|
}
|
|
}()
|
|
}
|
|
|
|
// PropagateCancel creates a derived context that is cancelled when either the parent
|
|
// context is cancelled or the cancel channel is closed.
|
|
func (cm *CancelManager) PropagateCancel(parent context.Context, cancelCh <-chan struct{}) (context.Context, context.CancelFunc) {
|
|
ctx, cancel := context.WithCancel(parent)
|
|
go func() {
|
|
select {
|
|
case <-cancelCh:
|
|
cancel()
|
|
case <-ctx.Done():
|
|
}
|
|
}()
|
|
return ctx, cancel
|
|
}
|