Set keepalive policy for prism grpc server.

Author: shunping

sdks/go/pkg/beam/runners/prism/internal/jobservices/server.go

@@ -30,6 +30,7 @@ import (
 	jobpb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/jobmanagement_v1"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/runners/prism/internal/worker"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/keepalive"
 )
 
 type Server struct {
@@ -78,8 +79,15 @@ func NewServer(port int, execute func(*Job)) *Server {
 		logger:  slog.Default(), // TODO substitute with a configured logger.
 	}
 	s.logger.Info("Serving JobManagement", slog.String("endpoint", s.Endpoint()))
+
+	// Define the keepalive enforcement policy
+	ep := keepalive.EnforcementPolicy{
+		MinTime:             20 * time.Second, // Minimum duration a client should wait before sending a keepalive ping
+		PermitWithoutStream: true,             // Allow pings even if there are no active streams
+	}
 	opts := []grpc.ServerOption{
 		grpc.MaxRecvMsgSize(math.MaxInt32),
+		grpc.KeepaliveEnforcementPolicy(ep),
 	}
 	s.server = grpc.NewServer(opts...)
 	jobpb.RegisterJobServiceServer(s.server, s)