Configure task proxy timeouts

README.md

@@ -15,6 +15,46 @@ For more information, refer to:
 - [Spark UI](https://ytsaurus.tech/docs/user-guide/data-processing/spyt/spark-ui) to learn how to open UI of [SPYT](https://ytsaurus.tech/docs/en/user-guide/data-processing/spyt/overview) clusters and jobs,
 - [Admin docs](https://ytsaurus.tech/docs/admin-guide/install-task-proxy) for installation instructions.
 
+## Annotating an operation
+
+To publish services from a regular YTsaurus operation, add the `task_proxy` annotation to its specification. `enabled` is required; `tasks_info` describes services by task name, service name, protocol, and zero-based job port index.
+
+```yson
+<"task_proxy"={
+    "enabled"=%true;
+    "tasks_info"={
+        "worker"={
+            "api"={
+                "protocol"="http";
+                "port_index"=0;
+            };
+            "grpc"={
+                "protocol"="grpc";
+                "port_index"=1;
+            };
+        };
+    };
+}>
+```
+
+`protocol` must be `http` or `grpc`. If `tasks_info` is omitted, task-proxy publishes every job port as an HTTP service named `port<N>`.
+
+The annotation can also override request timeouts for every service in that operation:
+
+```yson
+<"task_proxy"={
+    "enabled"=%true;
+    "route_timeout_seconds"=600;
+    "stream_idle_timeout_seconds"=120;
+}>
+```
+
+- `route_timeout_seconds` is the maximum time to receive a complete upstream response after Envoy has received the full request.
+- `stream_idle_timeout_seconds` is the maximum period without request or response traffic.
+- Both values are non-negative integer seconds. `0` explicitly disables the corresponding timeout; an omitted value inherits the global Helm setting.
+
+The Helm defaults are 2 seconds for connecting to a job, 15 seconds for a complete response, and 300 seconds for a stream idle period. Configure them through `timeouts.connectTimeoutSeconds`, `timeouts.routeTimeoutSeconds`, and `timeouts.streamIdleTimeoutSeconds`.
+
 ## Development
 
 Install chart to cluster from local directory using:

chart/templates/deployment.yaml

@@ -54,6 +54,9 @@ spec:
         - "-base-domain={{ .Values.baseDomain }}"
         - "-dir-path={{ .Values.dirPath }}"
         - "-discovery-period-seconds={{ .Values.discoveryPeriodSeconds }}"
+        - "-connect-timeout-seconds={{ .Values.timeouts.connectTimeoutSeconds }}"
+        - "-route-timeout-seconds={{ .Values.timeouts.routeTimeoutSeconds }}"
+        - "-stream-idle-timeout-seconds={{ .Values.timeouts.streamIdleTimeoutSeconds }}"
         - "-auth-enabled={{ .Values.auth.enabled }}"
         - "-auth-cookie-name={{ .Values.auth.cookieName }}"
         - "-auth-cache-enabled={{ .Values.auth.cache.enabled }}"

chart/values.yaml

@@ -12,6 +12,14 @@ dirPath: //sys/task_proxies
 
 discoveryPeriodSeconds: 60
 
+timeouts:
+  # Maximum time to establish a TCP connection to a task job.
+  connectTimeoutSeconds: 2
+  # Maximum time to wait for a complete upstream response. Set to 0 to disable.
+  routeTimeoutSeconds: 15
+  # Maximum period without request or response traffic. Set to 0 to disable.
+  streamIdleTimeoutSeconds: 300
+
 auth:
   enabled: true
   cookieName: YTCypressCookie

server/main.go

@@ -19,20 +19,24 @@ import (
 
 func main() {
 	ctx := context.Background()
+	defaultTimeoutConfig := pkg.DefaultTaskProxyTimeoutConfig()
 
 	var args struct {
-		ytProxy                 string
-		ytTokenPath             string
-		baseDomain              string
-		dirPath                 string
-		discoveryPeriodSeconds  uint
-		authEnabled             bool
-		authCookieName          string
-		authCacheEnabled        bool
-		authCacheTTLSeconds     int
-		authCacheCapacity       int
-		authCacheMaxConcurrency int
-		authCacheRefreshBefore  int
+		ytProxy                  string
+		ytTokenPath              string
+		baseDomain               string
+		dirPath                  string
+		discoveryPeriodSeconds   uint
+		authEnabled              bool
+		authCookieName           string
+		authCacheEnabled         bool
+		authCacheTTLSeconds      int
+		authCacheCapacity        int
+		authCacheMaxConcurrency  int
+		authCacheRefreshBefore   int
+		connectTimeoutSeconds    int
+		routeTimeoutSeconds      int
+		streamIdleTimeoutSeconds int
 	}
 	flag.StringVar(&args.ytProxy, "yt-proxy", "", "YT proxy host")
 	flag.StringVar(&args.ytTokenPath, "yt-token-path", "", "YT token path")
@@ -46,6 +50,9 @@ func main() {
 	flag.IntVar(&args.authCacheCapacity, "auth-cache-capacity", 0, "auth cache maximum number of entries (0 means unlimited)")
 	flag.IntVar(&args.authCacheMaxConcurrency, "auth-cache-max-concurrent-backend-requests", 0, "auth cache max concurrent backend requests per key on misses (0 means unlimited)")
 	flag.IntVar(&args.authCacheRefreshBefore, "auth-cache-refresh-before-seconds", 0, "auth cache proactive refresh threshold in seconds before TTL deadline (0 disables proactive refresh)")
+	flag.IntVar(&args.connectTimeoutSeconds, "connect-timeout-seconds", int(defaultTimeoutConfig.ConnectTimeout/time.Second), "maximum time in seconds to establish an upstream job connection")
+	flag.IntVar(&args.routeTimeoutSeconds, "route-timeout-seconds", int(defaultTimeoutConfig.RouteTimeout/time.Second), "maximum time in seconds to wait for a complete upstream response (0 disables the timeout)")
+	flag.IntVar(&args.streamIdleTimeoutSeconds, "stream-idle-timeout-seconds", int(defaultTimeoutConfig.StreamIdleTimeout/time.Second), "maximum idle time in seconds for an upstream request or response stream (0 disables the timeout)")
 	flag.Parse()
 
 	if args.ytProxy == "" {
@@ -75,6 +82,26 @@ func main() {
 	if args.authCacheRefreshBefore < 0 {
 		log.Fatal("'auth-cache-refresh-before-seconds' argument must be non-negative")
 	}
+	connectTimeout, err := pkg.DurationFromSeconds(args.connectTimeoutSeconds)
+	if err != nil {
+		log.Fatalf("invalid connect timeout: %v", err)
+	}
+	routeTimeout, err := pkg.DurationFromSeconds(args.routeTimeoutSeconds)
+	if err != nil {
+		log.Fatalf("invalid route timeout: %v", err)
+	}
+	streamIdleTimeout, err := pkg.DurationFromSeconds(args.streamIdleTimeoutSeconds)
+	if err != nil {
+		log.Fatalf("invalid stream idle timeout: %v", err)
+	}
+	timeoutConfig := pkg.TaskProxyTimeoutConfig{
+		ConnectTimeout:    connectTimeout,
+		RouteTimeout:      routeTimeout,
+		StreamIdleTimeout: streamIdleTimeout,
+	}
+	if err := timeoutConfig.Validate(); err != nil {
+		log.Fatalf("invalid task proxy timeout configuration: %v", err)
+	}
 
 	ytTokenBytes, err := os.ReadFile(args.ytTokenPath)
 	if err != nil {
@@ -109,7 +136,7 @@ func main() {
 		RefreshBeforeSeconds:         args.authCacheRefreshBefore,
 	})
 
-	taskUpdater := pkg.CreateTaskUpdater(args.baseDomain, tls, args.authEnabled, authServer, taskDiscovery, cache)
+	taskUpdater := pkg.CreateTaskUpdater(args.baseDomain, tls, args.authEnabled, timeoutConfig, authServer, taskDiscovery, cache)
 
 	go func() {
 		if err := pkg.ServeMetrics(pkg.DefaultGatherer()); err != nil {

server/pkg/discovery.go

@@ -3,6 +3,7 @@ package pkg
 import (
 	"context"
 	"fmt"
+	"math"
 	"net"
 	"net/url"
 	"strconv"
@@ -14,7 +15,17 @@ import (
 	ytsdk "go.ytsaurus.tech/yt/go/yt"
 )
 
-const servicesTableName = "services"
+const (
+	servicesTableName = "services"
+
+	taskProxyAnnotationKey          = "task_proxy"
+	taskProxyEnabledKey             = "enabled"
+	taskProxyTasksInfoKey           = "tasks_info"
+	taskProxyProtocolKey            = "protocol"
+	taskProxyPortIndexKey           = "port_index"
+	taskProxyRouteTimeoutSecondsKey = "route_timeout_seconds"
+	taskProxyStreamIdleTimeoutKey   = "stream_idle_timeout_seconds"
+)
 
 type taskDiscovery struct {
 	baseDomain string
@@ -63,7 +74,7 @@ func (d *taskDiscovery) Discovery(ctx context.Context) (TaskList, error) {
 				d.logger.Errorf("unable to process SPYT standalone cluster operation %q: %v", op.ID, err)
 				continue
 			}
-		} else if _, ok := annotations["task_proxy"]; ok {
+		} else if _, ok := annotations[taskProxyAnnotationKey]; ok {
 			opTasks, err = d.processTaskProxyAnnotatedOperation(ctx, op)
 			if err != nil {
 				d.logger.Errorf("unable to process task proxy annotated operation %q: %v", op.ID, err)
@@ -195,8 +206,8 @@ func (d *taskDiscovery) processSPYTStandaloneClusterOperation(ctx context.Contex
 }
 
 func (d *taskDiscovery) processTaskProxyAnnotatedOperation(ctx context.Context, op ytsdk.OperationStatus) ([]Task, error) {
-	taskProxyAnnotation := op.RuntimeParameters.Annotations["task_proxy"]
-	taskServiceInfos := parseTaskProxyAnnotation(taskProxyAnnotation)
+	taskProxyAnnotation := op.RuntimeParameters.Annotations[taskProxyAnnotationKey]
+	taskServiceInfos, timeoutOverrides := parseTaskProxyAnnotation(taskProxyAnnotation)
 	if taskServiceInfos == nil {
 		return nil, fmt.Errorf("invalid task_proxy annotation: %v", taskProxyAnnotation)
 	}
@@ -250,11 +261,12 @@ func (d *taskDiscovery) processTaskProxyAnnotatedOperation(ctx context.Context,
 			hostParts := strings.Split(job.Address, ":") // job address contains port also
 
 			taskProto := Task{
-				operationID:    op.ID.String(),
-				operationAlias: parseOperationAlias(op),
-				taskName:       job.TaskName,
-				service:        serviceInfo.service,
-				protocol:       serviceInfo.protocol,
+				operationID:      op.ID.String(),
+				operationAlias:   parseOperationAlias(op),
+				taskName:         job.TaskName,
+				service:          serviceInfo.service,
+				protocol:         serviceInfo.protocol,
+				timeoutOverrides: timeoutOverrides,
 			}
 			if _, ok := idToTask[taskProto.ID()]; !ok {
 				idToTask[taskProto.ID()] = &taskProto
@@ -371,32 +383,36 @@ type taskServiceInfo struct {
 	portIndex int
 }
 
-func parseTaskProxyAnnotation(taskProxyAny any) []taskServiceInfo {
+func parseTaskProxyAnnotation(taskProxyAny any) ([]taskServiceInfo, TaskTimeoutOverrides) {
 	taskProxy, ok := taskProxyAny.(map[string]any)
 	if !ok {
-		return nil
+		return nil, TaskTimeoutOverrides{}
 	}
-	enabledAny, ok := taskProxy["enabled"]
+	enabledAny, ok := taskProxy[taskProxyEnabledKey]
 	if !ok {
-		return nil
+		return nil, TaskTimeoutOverrides{}
 	}
 	enabled, ok := enabledAny.(bool)
 	if !ok {
-		return nil
+		return nil, TaskTimeoutOverrides{}
 	}
 	if !enabled {
-		return nil
+		return nil, TaskTimeoutOverrides{}
+	}
+	timeoutOverrides, ok := parseTaskTimeoutOverrides(taskProxy)
+	if !ok {
+		return nil, TaskTimeoutOverrides{}
 	}
 
 	taskServiceInfos := make([]taskServiceInfo, 0)
-	tasksInfoAny, ok := taskProxy["tasks_info"]
+	tasksInfoAny, ok := taskProxy[taskProxyTasksInfoKey]
 	if !ok {
-		return taskServiceInfos
+		return taskServiceInfos, timeoutOverrides
 	}
 
 	tasksInfo, ok := tasksInfoAny.(map[string]any)
 	if !ok {
-		return taskServiceInfos
+		return taskServiceInfos, timeoutOverrides
 	}
 
 	for task, infoAny := range tasksInfo {
@@ -409,7 +425,7 @@ func parseTaskProxyAnnotation(taskProxyAny any) []taskServiceInfo {
 			if !ok {
 				continue
 			}
-			protocolAny, ok := info["protocol"]
+			protocolAny, ok := info[taskProxyProtocolKey]
 			if !ok {
 				continue
 			}
@@ -420,31 +436,16 @@ func parseTaskProxyAnnotation(taskProxyAny any) []taskServiceInfo {
 			if protocol != string(HTTP) && protocol != string(GRPC) {
 				continue
 			}
-			portIndexAny, ok := info["port_index"]
+			portIndexAny, ok := info[taskProxyPortIndexKey]
 			if !ok {
 				continue
 			}
-			var portIndex int
-			switch v := portIndexAny.(type) {
-			case int:
-				portIndex = v
-			case int64:
-				portIndex = int(v)
-			case int32:
-				portIndex = int(v)
-			case int16:
-				portIndex = int(v)
-			case int8:
-				portIndex = int(v)
-			case uint64:
-				portIndex = int(v)
-			case uint32:
-				portIndex = int(v)
-			case uint16:
-				portIndex = int(v)
-			case uint8:
-				portIndex = int(v)
-			default:
+			portIndex64, ok := parseInteger(portIndexAny)
+			if !ok {
+				continue
+			}
+			portIndex := int(portIndex64)
+			if int64(portIndex) != portIndex64 {
 				continue
 			}
 			taskServiceInfos = append(taskServiceInfos, taskServiceInfo{
@@ -456,7 +457,62 @@ func parseTaskProxyAnnotation(taskProxyAny any) []taskServiceInfo {
 		}
 	}
 
-	return taskServiceInfos
+	return taskServiceInfos, timeoutOverrides
+}
+
+func parseTaskTimeoutOverrides(taskProxy map[string]any) (TaskTimeoutOverrides, bool) {
+	var overrides TaskTimeoutOverrides
+	if value, ok := taskProxy[taskProxyRouteTimeoutSecondsKey]; ok {
+		timeout, ok := parseTimeoutSeconds(value)
+		if !ok {
+			return TaskTimeoutOverrides{}, false
+		}
+		overrides.routeTimeout = &timeout
+	}
+	if value, ok := taskProxy[taskProxyStreamIdleTimeoutKey]; ok {
+		timeout, ok := parseTimeoutSeconds(value)
+		if !ok {
+			return TaskTimeoutOverrides{}, false
+		}
+		overrides.streamIdleTimeout = &timeout
+	}
+	return overrides, true
+}
+
+func parseTimeoutSeconds(value any) (time.Duration, bool) {
+	seconds, ok := parseInteger(value)
+	if !ok || seconds < 0 || seconds > int64(math.MaxInt64/time.Second) {
+		return 0, false
+	}
+	return time.Duration(seconds) * time.Second, true
+}
+
+func parseInteger(value any) (int64, bool) {
+	switch v := value.(type) {
+	case int:
+		return int64(v), true
+	case int64:
+		return v, true
+	case int32:
+		return int64(v), true
+	case int16:
+		return int64(v), true
+	case int8:
+		return int64(v), true
+	case uint64:
+		if v > uint64(math.MaxInt64) {
+			return 0, false
+		}
+		return int64(v), true
+	case uint32:
+		return int64(v), true
+	case uint16:
+		return int64(v), true
+	case uint8:
+		return int64(v), true
+	default:
+		return 0, false
+	}
 }
 
 func makeHostPortFromNode(node string) (*HostPort, error) {