[Feature] Add timeout for apiserver grpc server

apiserver/cmd/main.go

@@ -11,6 +11,7 @@ import (
 	"path"
 	"strings"
 	"sync/atomic"
+	"time"
 
 	assetfs "github.com/elazarl/go-bindata-assetfs"
 	grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware"
@@ -50,11 +51,21 @@ func main() {
 		_ = flagSet.Set("log_file", *logFile)
 	}
 
+	grpcTimeout := 60 * time.Second // Default timeout
+	if timeoutStr := os.Getenv("GRPC_SERVER_TIMEOUT"); timeoutStr != "" {
+		if timeout, err := time.ParseDuration(timeoutStr); err == nil {
+			grpcTimeout = timeout
+			klog.Infof("gRPC servier timeout set to %v", grpcTimeout)
+		} else {
+			klog.Warningf("Invalid GRPC_SERVER_TIMEOUT value: %v, using default timeout (60 seconds)", err)
+		}
+	}
+
 	clientManager := manager.NewClientManager()
 	resourceManager := manager.NewResourceManager(&clientManager)
 
 	atomic.StoreInt32(&healthy, 1)
-	go startRPCServer(resourceManager)
+	go startRPCServer(resourceManager, grpcTimeout)
 	startHttpProxy()
 	// See also https://gist.github.com/enricofoltran/10b4a980cd07cb02836f70a4ab3e72d7
 	quit := make(chan os.Signal, 1)
@@ -70,7 +81,7 @@ func main() {
 
 type RegisterHttpHandlerFromEndpoint func(ctx context.Context, mux *runtime.ServeMux, endpoint string, opts []grpc.DialOption) error
 
-func startRPCServer(resourceManager *manager.ResourceManager) {
+func startRPCServer(resourceManager *manager.ResourceManager, grpcTimeout time.Duration) {
 	klog.Infof("Starting gRPC server at port %s", *rpcPortFlag)
 
 	listener, err := net.Listen("tcp", *rpcPortFlag)
@@ -86,8 +97,13 @@ func startRPCServer(resourceManager *manager.ResourceManager) {
 
 	s := grpc.NewServer(
 		grpc.StreamInterceptor(grpc_prometheus.StreamServerInterceptor),
-		grpc.UnaryInterceptor(grpc_middleware.ChainUnaryServer(grpc_prometheus.UnaryServerInterceptor, interceptor.APIServerInterceptor)),
-		grpc.MaxRecvMsgSize(math.MaxInt32))
+		grpc.UnaryInterceptor(grpc_middleware.ChainUnaryServer(
+			interceptor.TimeoutInterceptor(grpcTimeout),
+			grpc_prometheus.UnaryServerInterceptor,
+			interceptor.APIServerInterceptor,
+		)),
+		grpc.MaxRecvMsgSize(math.MaxInt32),
+	)
 	api.RegisterClusterServiceServer(s, clusterServer)
 	api.RegisterComputeTemplateServiceServer(s, templateServer)
 	api.RegisterRayJobServiceServer(s, jobServer)

apiserver/pkg/interceptor/interceptor.go

@@ -2,6 +2,8 @@ package interceptor
 
 import (
 	"context"
+	"fmt"
+	"time"
 
 	"google.golang.org/grpc"
 	klog "k8s.io/klog/v2"
@@ -19,3 +21,41 @@ func APIServerInterceptor(ctx context.Context, req interface{}, info *grpc.Unary
 	klog.Infof("%v handler finished", info.FullMethod)
 	return
 }
+
+// TimeoutInterceptor implements UnaryServerInterceptor that sets the timeout for the request
+func TimeoutInterceptor(timeout time.Duration) grpc.UnaryServerInterceptor {
+	return func(
+		ctx context.Context,
+		req interface{},
+		_ *grpc.UnaryServerInfo,
+		handler grpc.UnaryHandler,
+	) (interface{}, error) {
+		// Create a context with timeout
+		ctx, cancel := context.WithTimeout(ctx, timeout)
+		defer cancel()
+
+		// Channel to capture execution result
+		done := make(chan struct{})
+		var (
+			resp interface{}
+			err  error
+		)
+
+		go func() {
+			resp, err = handler(ctx, req)
+			close(done)
+		}()
+
+		select {
+		case <-ctx.Done():
+			// Raise error if time out
+			if ctx.Err() == context.DeadlineExceeded {
+				return nil, fmt.Errorf("grpc server timed out")
+			}
+			return nil, ctx.Err()
+		case <-done:
+			// Handler finished
+			return resp, err
+		}
+	}
+}

apiserver/pkg/interceptor/interceptor_test.go

@@ -4,9 +4,11 @@ import (
 	"bytes"
 	"context"
 	"errors"
+	"fmt"
 	"io"
 	"os"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -192,3 +194,63 @@ func TestAPIServerInterceptorLogging(t *testing.T) {
 		})
 	}
 }
+
+func TestTimeoutInterceptor(t *testing.T) {
+	tests := []struct {
+		expectedError  error
+		name           string
+		timeout        time.Duration
+		handlerDelay   time.Duration
+		expectedCalled bool
+	}{
+		{
+			name:           "handler completes before timeout",
+			timeout:        100 * time.Millisecond,
+			handlerDelay:   50 * time.Millisecond,
+			expectedError:  nil,
+			expectedCalled: true,
+		},
+		{
+			name:           "handler exceeds timeout",
+			timeout:        50 * time.Millisecond,
+			handlerDelay:   100 * time.Millisecond,
+			expectedError:  fmt.Errorf("grpc server timed out"),
+			expectedCalled: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Create test context and request
+			ctx := context.Background()
+			req := "test_request"
+			handler := &mockHandler{}
+
+			// Create the interceptor with the specified timeout
+			interceptor := TimeoutInterceptor(tt.timeout)
+
+			// Call the interceptor
+			resp, err := interceptor(
+				ctx,
+				req,
+				&grpc.UnaryServerInfo{FullMethod: "TestTimeoutMethod"},
+				func(ctx context.Context, req interface{}) (interface{}, error) {
+					time.Sleep(tt.handlerDelay)
+					return handler.Handle(ctx, req)
+				},
+			)
+
+			// Verify response and error
+			if tt.expectedError == nil {
+				// Verify handler was called
+				assert.Equal(t, tt.expectedCalled, handler.called, "handler call status should match expected")
+
+				require.NoError(t, err)
+				assert.Equal(t, "test_response", resp, "response should match expected")
+			} else {
+				require.Error(t, err)
+				require.EqualError(t, err, tt.expectedError.Error(), "A matching error is expected")
+			}
+		})
+	}
+}