From 416b6e0f31855ce3947e148d800de3b0505648fc Mon Sep 17 00:00:00 2001
From: "Lucian I. Last" <li@last.nl>
Date: Thu, 19 Dec 2024 15:26:44 +0100
Subject: [PATCH] add GRPC_HIGH_SCALE=true

---
 .example.env     |  1 +
 config/config.go |  2 ++
 grpc/server.go   | 18 +++++++++++++++++-
 3 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/.example.env b/.example.env
index 318c9c4..7cced6c 100644
--- a/.example.env
+++ b/.example.env
@@ -9,6 +9,7 @@ TB_ADDRESSES=3033
 # USE_GRPC=true
 # GRPC_HEALTH_SERVER=true
 # GRPC_REFLECTION=true
+# GRPC_HIGH_SCALE=true
 
 # IS_BUFFERED=true
 # BUFFER_SIZE=20
diff --git a/config/config.go b/config/config.go
index a4ef18f..7e29277 100644
--- a/config/config.go
+++ b/config/config.go
@@ -24,6 +24,7 @@ type config struct {
 	UseGrpc          bool
 	GrpcHealthServer bool
 	GrpcReflection   bool
+	GrpcHighScale    bool
 
 	IsBuffered    bool
 	BufferSize    int
@@ -96,6 +97,7 @@ func NewConfig() (ok bool) {
 		UseGrpc:          useGrpc,
 		GrpcHealthServer: os.Getenv("GRPC_HEALTH_SERVER") == "true",
 		GrpcReflection:   os.Getenv("GRPC_REFLECTION") == "true",
+		GrpcHighScale:    os.Getenv("GRPC_HIGH_SCALE") == "true",
 
 		IsBuffered:    isBuffered,
 		BufferSize:    bufferSize,
diff --git a/grpc/server.go b/grpc/server.go
index d55bcd6..805b64b 100644
--- a/grpc/server.go
+++ b/grpc/server.go
@@ -5,6 +5,7 @@ import (
 	"log/slog"
 	"net"
 	"os"
+	"time"
 
 	"github.com/lil5/tigerbeetle_api/config"
 	"github.com/lil5/tigerbeetle_api/metrics"
@@ -13,6 +14,7 @@ import (
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/health"
 	healthpb "google.golang.org/grpc/health/grpc_health_v1"
+	"google.golang.org/grpc/keepalive"
 	"google.golang.org/grpc/reflection"
 
 	grpcprom "github.com/grpc-ecosystem/go-grpc-middleware/providers/prometheus"
@@ -28,7 +30,21 @@ func NewServer() {
 		slog.Error("Failed to listen", "error", err)
 		os.Exit(1)
 	}
-	s := grpc.NewServer()
+
+	var s *grpc.Server
+	if config.Config.GrpcHighScale {
+		s = grpc.NewServer(
+			grpc.KeepaliveParams(keepalive.ServerParameters{
+				MaxConnectionAge: 5 * time.Minute,
+				Time:             60 * time.Second,
+				Timeout:          20 * time.Second,
+			}),
+			grpc.MaxConcurrentStreams(50_000),
+		)
+	} else {
+		s = grpc.NewServer()
+	}
+
 	app := NewApp()
 	defer app.Close()
 	proto.RegisterTigerBeetleServer(s, app)