Add configurable timeout for pending backend dials

cmd/server/app/options/options.go

@@ -122,6 +122,8 @@ type ProxyRunOptions struct {
 	NeedsKubernetesClient bool
 	// Graceful shutdown timeout duration
 	GracefulShutdownTimeout time.Duration
+	// Backend dial timeout duration for requests from the server to backend agents.
+	BackendDialTimeout time.Duration
 }
 
 func (o *ProxyRunOptions) Flags() *pflag.FlagSet {
@@ -164,6 +166,7 @@ func (o *ProxyRunOptions) Flags() *pflag.FlagSet {
 	flags.StringVar(&o.LeaseNamespace, "lease-namespace", o.LeaseNamespace, "The namespace where lease objects are managed by the controller.")
 	flags.StringVar(&o.LeaseLabel, "lease-label", o.LeaseLabel, "The labels on which the lease objects are managed.")
 	flags.DurationVar(&o.GracefulShutdownTimeout, "graceful-shutdown-timeout", o.GracefulShutdownTimeout, "Timeout duration for graceful shutdown of the server. The server will wait for active connections to close before forcefully terminating. Set to 0 to disable graceful shutdown (default: 0).")
+	flags.DurationVar(&o.BackendDialTimeout, "backend-dial-timeout", o.BackendDialTimeout, "Timeout duration for sending DIAL_REQ packets to backend agent streams and waiting for DIAL_RSP. Set to 0 to disable timeout (default: 0).")
 	flags.Bool("warn-on-channel-limit", true, "This behavior is now thread safe and always on. This flag will be removed in a future release.")
 	flags.MarkDeprecated("warn-on-channel-limit", "This behavior is now thread safe and always on. This flag will be removed in a future release.")
 
@@ -209,6 +212,7 @@ func (o *ProxyRunOptions) Print() {
 	klog.V(1).Infof("TLSMinVersion set to %q.\n", o.TLSMinVersion)
 	klog.V(1).Infof("XfrChannelSize set to %d.\n", o.XfrChannelSize)
 	klog.V(1).Infof("GracefulShutdownTimeout set to %v.\n", o.GracefulShutdownTimeout)
+	klog.V(1).Infof("BackendDialTimeout set to %v.\n", o.BackendDialTimeout)
 }
 
 func (o *ProxyRunOptions) Validate() error {
@@ -368,6 +372,9 @@ func (o *ProxyRunOptions) Validate() error {
 	if o.GracefulShutdownTimeout < 0 {
 		return fmt.Errorf("graceful-shutdown-timeout must be >= 0, got %v", o.GracefulShutdownTimeout)
 	}
+	if o.BackendDialTimeout < 0 {
+		return fmt.Errorf("backend-dial-timeout must be >= 0, got %v", o.BackendDialTimeout)
+	}
 
 	o.NeedsKubernetesClient = usingServiceAccountAuth || o.EnableLeaseController
 
@@ -414,6 +421,7 @@ func NewProxyRunOptions() *ProxyRunOptions {
 		LeaseNamespace:            "kube-system",
 		LeaseLabel:                "k8s-app=konnectivity-server",
 		GracefulShutdownTimeout:   0,
+		BackendDialTimeout:        0,
 	}
 	return &o
 }

cmd/server/app/options/options_test.go

@@ -66,6 +66,7 @@ func TestDefaultServerOptions(t *testing.T) {
 	assertDefaultValue(t, "XfrChannelSize", defaultServerOptions.XfrChannelSize, 10)
 	assertDefaultValue(t, "APIContentType", defaultServerOptions.APIContentType, "application/vnd.kubernetes.protobuf")
 	assertDefaultValue(t, "GracefulShutdownTimeout", defaultServerOptions.GracefulShutdownTimeout, 0*time.Second)
+	assertDefaultValue(t, "BackendDialTimeout", defaultServerOptions.BackendDialTimeout, 0*time.Second)
 
 }
 
@@ -220,6 +221,21 @@ func TestValidate(t *testing.T) {
 			value:    30 * time.Second,
 			expected: nil,
 		},
+		"NegativeBackendDialTimeout": {
+			field:    "BackendDialTimeout",
+			value:    -1 * time.Second,
+			expected: fmt.Errorf("backend-dial-timeout must be >= 0, got -1s"),
+		},
+		"ZeroBackendDialTimeout": {
+			field:    "BackendDialTimeout",
+			value:    0 * time.Second,
+			expected: nil,
+		},
+		"PositiveBackendDialTimeout": {
+			field:    "BackendDialTimeout",
+			value:    30 * time.Second,
+			expected: nil,
+		},
 	} {
 		t.Run(desc, func(t *testing.T) {
 			testServerOptions := NewProxyRunOptions()

cmd/server/app/server.go

@@ -140,6 +140,7 @@ func (p *Proxy) Run(o *options.ProxyRunOptions, stopCh <-chan struct{}) error {
 		return err
 	}
 	p.server = server.NewProxyServer(o.ServerID, ps, o.ServerCount, authOpt, o.XfrChannelSize)
+	p.server.SetBackendDialTimeout(o.BackendDialTimeout)
 
 	frontendStop, err := p.runFrontendServer(ctx, o, p.server)
 	if err != nil {

pkg/server/backend_manager.go

@@ -43,14 +43,17 @@ import (
 // agent.AgentService_ConnectServer, provides synchronization and
 // emits common stream metrics.
 type Backend struct {
-	sendLock sync.Mutex
-	recvLock sync.Mutex
-	conn     agent.AgentService_ConnectServer
+	sendLock   sync.Mutex
+	recvLock   sync.Mutex
+	retireOnce sync.Once
+	conn       agent.AgentService_ConnectServer
 
 	// cached from conn.Context()
 	id     string
 	idents header.Identifiers
 
+	done chan struct{}
+
 	// draining indicates if this backend is draining and should not accept new connections
 	draining atomic.Bool
 }
@@ -65,6 +68,17 @@ func (b *Backend) SetDraining() {
 	b.draining.Store(true)
 }
 
+func (b *Backend) Retire() {
+	b.SetDraining()
+	b.retireOnce.Do(func() {
+		close(b.done)
+	})
+}
+
+func (b *Backend) Done() <-chan struct{} {
+	return b.done
+}
+
 func (b *Backend) Send(p *client.Packet) error {
 	b.sendLock.Lock()
 	defer b.sendLock.Unlock()
@@ -145,7 +159,7 @@ func NewBackend(conn agent.AgentService_ConnectServer) (*Backend, error) {
 	if err != nil {
 		return nil, err
 	}
-	return &Backend{conn: conn, id: agentID, idents: agentIdentifiers}, nil
+	return &Backend{conn: conn, id: agentID, idents: agentIdentifiers, done: make(chan struct{})}, nil
 }
 
 // BackendStorage is an interface to manage the storage of the backend

pkg/server/metrics/metrics.go

@@ -328,6 +328,7 @@ const (
 	DialFailureUnrecognizedResponse DialFailureReason = "unrecognized_response" // Dial repsonse received for unrecognozide dial ID.
 	DialFailureSendResponse         DialFailureReason = "send_rsp"              // Successful dial response from agent, but failed to send to frontend.
 	DialFailureBackendClose         DialFailureReason = "backend_close"         // Received a DIAL_CLS from the backend before the dial completed.
+	DialFailureBackendDialTimeout   DialFailureReason = "backend_dial_timeout"  // Timed out sending DIAL_REQ to or waiting for DIAL_RSP from the backend.
 	DialFailureFrontendClose        DialFailureReason = "frontend_close"        // Received a DIAL_CLS from the frontend before the dial completed.
 )
 

pkg/server/server.go

@@ -90,6 +90,10 @@ const (
 	ModeHTTPConnect = "http-connect"
 )
 
+const defaultBackendDialTimeout = 0
+
+var errBackendDialTimeout = errors.New("timed out waiting for backend dial")
+
 type ProxyClientConnection struct {
 	Mode        string
 	HTTP        io.ReadWriter
@@ -274,6 +278,8 @@ type ProxyServer struct {
 	// TODO: move strategies into BackendStorage
 	proxyStrategies []proxystrategies.ProxyStrategy
 	xfrChannelSize  int
+
+	backendDialTimeout time.Duration
 }
 
 // AgentTokenAuthenticationOptions contains list of parameters required for agent token based authentication
@@ -317,6 +323,65 @@ func (s *ProxyServer) getBackend(reqHost string) (*Backend, error) {
 	return nil, &ErrNotFound{}
 }
 
+func (s *ProxyServer) sendDialRequestToBackend(backend *Backend, pkt *client.Packet) error {
+	timeout := s.backendDialTimeout
+	if timeout <= 0 {
+		return backend.Send(pkt)
+	}
+
+	errCh := make(chan error, 1)
+	go func() {
+		errCh <- backend.Send(pkt)
+	}()
+
+	timer := time.NewTimer(timeout)
+	defer timer.Stop()
+
+	ctx := backend.Context()
+	select {
+	case err := <-errCh:
+		return err
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-timer.C:
+		s.retireBackend(backend, "backend dial request send timed out")
+		return errBackendDialTimeout
+	}
+}
+
+func (s *ProxyServer) startPendingDialTimeout(random int64, backend *Backend, frontend *GrpcFrontend) {
+	timeout := s.backendDialTimeout
+	if timeout <= 0 {
+		return
+	}
+
+	go func() {
+		timer := time.NewTimer(timeout)
+		defer timer.Stop()
+
+		select {
+		case <-timer.C:
+			if s.PendingDial.Remove(random) == nil {
+				return
+			}
+			metrics.Metrics.ObserveDialFailure(metrics.DialFailureBackendDialTimeout)
+			resp := &client.Packet{
+				Type: client.PacketType_DIAL_RSP,
+				Payload: &client.Packet_DialResponse{
+					DialResponse: &client.DialResponse{
+						Random: random,
+						Error:  errBackendDialTimeout.Error(),
+					},
+				},
+			}
+			if err := frontend.Send(resp); err != nil {
+				klog.V(5).InfoS("Failed to send DIAL_RSP for backend dial timeout", "error", err, "dialID", random)
+			}
+		case <-backend.Context().Done():
+		}
+	}()
+}
+
 func (s *ProxyServer) addBackend(backend *Backend) {
 	// TODO: refactor BackendStorage to acquire lock once, not up to 3 times.
 	for _, bm := range s.BackendManagers {
@@ -330,6 +395,12 @@ func (s *ProxyServer) removeBackend(backend *Backend) {
 	}
 }
 
+func (s *ProxyServer) retireBackend(backend *Backend, reason string) {
+	backend.Retire()
+	klog.V(2).InfoS("Retire backend connection", "agentID", backend.GetAgentID(), "reason", reason)
+	s.removeBackend(backend)
+}
+
 func (s *ProxyServer) addEstablished(agentID string, connID int64, p *ProxyClientConnection) {
 	s.fmu.Lock()
 	defer s.fmu.Unlock()
@@ -455,12 +526,17 @@ func NewProxyServer(serverID string, proxyStrategies []proxystrategies.ProxyStra
 		BackendManagers:            bms,
 		AgentAuthenticationOptions: agentAuthenticationOptions,
 		// use the first backend-manager as the Readiness Manager
-		Readiness:       bms[0],
-		proxyStrategies: proxyStrategies,
-		xfrChannelSize:  channelSize,
+		Readiness:          bms[0],
+		proxyStrategies:    proxyStrategies,
+		xfrChannelSize:     channelSize,
+		backendDialTimeout: defaultBackendDialTimeout,
 	}
 }
 
+func (s *ProxyServer) SetBackendDialTimeout(timeout time.Duration) {
+	s.backendDialTimeout = timeout
+}
+
 // Proxy handles incoming streams from gRPC frontend.
 func (s *ProxyServer) Proxy(stream client.ProxyService_ProxyServer) error {
 	metrics.Metrics.ConnectionInc(metrics.Proxy)
@@ -613,11 +689,31 @@ func (s *ProxyServer) serveRecvFrontend(frontend *GrpcFrontend, recvCh <-chan *c
 					backend:     backend,
 					dialAddress: address,
 				})
-			if err := backend.Send(pkt); err != nil {
+			if err := s.sendDialRequestToBackend(backend, pkt); err != nil {
 				klog.ErrorS(err, "DIAL_REQ to Backend failed", "dialID", random)
-			} else {
-				klog.V(5).InfoS("DIAL_REQ sent to backend", "dialID", random)
+				if s.PendingDial.Remove(random) != nil {
+					reason := metrics.DialFailureBackendClose
+					if errors.Is(err, errBackendDialTimeout) {
+						reason = metrics.DialFailureBackendDialTimeout
+					}
+					metrics.Metrics.ObserveDialFailure(reason)
+				}
+				resp := &client.Packet{
+					Type: client.PacketType_DIAL_RSP,
+					Payload: &client.Packet_DialResponse{
+						DialResponse: &client.DialResponse{
+							Random: random,
+							Error:  err.Error(),
+						},
+					},
+				}
+				if err := frontend.Send(resp); err != nil {
+					klog.V(5).InfoS("Failed to send DIAL_RSP for backend send failure", "error", err, "dialID", random)
+				}
+				return
 			}
+			klog.V(5).InfoS("DIAL_REQ sent to backend", "dialID", random)
+			s.startPendingDialTimeout(random, backend, frontend)
 
 		case client.PacketType_CLOSE_REQ:
 			connID := pkt.GetCloseRequest().ConnectID
@@ -806,17 +902,21 @@ func (s *ProxyServer) Connect(stream agent.AgentService_ConnectServer) error {
 
 	go runpprof.Do(context.Background(), labels, func(context.Context) { s.serveRecvBackend(backend, agentID, recvCh) })
 
-	defer func() {
-		close(recvCh)
-	}()
-
-	stopCh := make(chan error)
+	stopCh := make(chan error, 1)
 	go runpprof.Do(context.Background(), labels, func(context.Context) { s.readBackendToChannel(backend, recvCh, stopCh) })
 
-	return <-stopCh
+	select {
+	case err := <-stopCh:
+		return err
+	case <-backend.Done():
+		klog.V(2).InfoS("Backend connection retired", "agentID", agentID)
+		return nil
+	}
 }
 
 func (s *ProxyServer) readBackendToChannel(backend *Backend, recvCh chan *client.Packet, stopCh chan error) {
+	defer close(recvCh)
+
 	agentID := backend.GetAgentID()
 	for {
 		in, err := backend.Recv()