[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"
 	"github.com/grpc-ecosystem/grpc-gateway/v2/runtime"
@@ -86,7 +87,7 @@ func startRPCServer(resourceManager *manager.ResourceManager) {
 
 	s := grpc.NewServer(
 		grpc.StreamInterceptor(grpc_prometheus.StreamServerInterceptor),
-		grpc.UnaryInterceptor(grpc_middleware.ChainUnaryServer(grpc_prometheus.UnaryServerInterceptor, interceptor.APIServerInterceptor)),
+		grpc.UnaryInterceptor(grpc_middleware.ChainUnaryServer(interceptor.TimeoutInterceptor(20*time.Second), grpc_prometheus.UnaryServerInterceptor, interceptor.APIServerInterceptor)),
 		grpc.MaxRecvMsgSize(math.MaxInt32))
 	api.RegisterClusterServiceServer(s, clusterServer)
 	api.RegisterComputeTemplateServiceServer(s, templateServer)

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,40 @@ func APIServerInterceptor(ctx context.Context, req interface{}, info *grpc.Unary
 	klog.Infof("%v handler finished", info.FullMethod)
 	return
 }
+
+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
+		}
+	}
+}