From e544cdb53d303a7282d8336468a086f0380cb3fe Mon Sep 17 00:00:00 2001
From: Hackerwins <susukang98@gmail.com>
Date: Sun, 21 Nov 2021 18:48:52 +0900
Subject: [PATCH] Add rpc-max-requests-bytes flag to set client request limit

---
 internal/cli/agent.go     |  6 ++++++
 test/helper/helper.go     |  4 +++-
 yorkie/config.go          |  9 +++++++--
 yorkie/config.sample.yml  |  9 ++++++---
 yorkie/rpc/config.go      | 12 ++++++++++--
 yorkie/rpc/server.go      |  2 ++
 yorkie/rpc/server_test.go |  3 ++-
 7 files changed, 36 insertions(+), 9 deletions(-)

diff --git a/internal/cli/agent.go b/internal/cli/agent.go
index f835855e4..fc4a23cf8 100644
--- a/internal/cli/agent.go
+++ b/internal/cli/agent.go
@@ -170,6 +170,12 @@ func init() {
 		"",
 		"RPC key file's path",
 	)
+	cmd.Flags().Uint64Var(
+		&conf.RPC.MaxRequestBytes,
+		"rpc-max-requests-bytes",
+		yorkie.DefaultRPCMaxRequestsBytes,
+		"Maximum client request size in bytes the server will accept.",
+	)
 	cmd.Flags().IntVar(
 		&conf.Profiling.Port,
 		"profiling-port",
diff --git a/test/helper/helper.go b/test/helper/helper.go
index 5a6e951d4..4c25b95c3 100644
--- a/test/helper/helper.go
+++ b/test/helper/helper.go
@@ -37,6 +37,7 @@ var testStartedAt int64
 // Below are the values of the Yorkie config used in the test.
 const (
 	RPCPort                    = 21101
+	RPCMaxRequestBytes         = 4 * 1024 * 1024
 	ProfilingPort              = 21102
 	MongoConnectionURI         = "mongodb://localhost:27017"
 	MongoConnectionTimeout     = "5s"
@@ -87,7 +88,8 @@ func TestConfig(authWebhook string) *yorkie.Config {
 	portOffset += 100
 	return &yorkie.Config{
 		RPC: &rpc.Config{
-			Port: RPCPort + portOffset,
+			Port:            RPCPort + portOffset,
+			MaxRequestBytes: RPCMaxRequestBytes,
 		},
 		Profiling: &profiling.Config{
 			Port: ProfilingPort + portOffset,
diff --git a/yorkie/config.go b/yorkie/config.go
index 91cd828cb..50bd182f9 100644
--- a/yorkie/config.go
+++ b/yorkie/config.go
@@ -34,8 +34,9 @@ import (
 
 // Below are the values of the default values of Yorkie config.
 const (
-	DefaultRPCPort       = 11101
-	DefaultProfilingPort = 11102
+	DefaultRPCPort             = 11101
+	DefaultProfilingPort       = 11102
+	DefaultRPCMaxRequestsBytes = 4 * 1024 * 1024 // 4MiB
 
 	DefaultMongoConnectionURI     = "mongodb://localhost:27017"
 	DefaultMongoConnectionTimeout = 5 * time.Second
@@ -120,6 +121,10 @@ func (c *Config) ensureDefaultValue() {
 		c.RPC.Port = DefaultRPCPort
 	}
 
+	if c.RPC.MaxRequestBytes == 0 {
+		c.RPC.MaxRequestBytes = DefaultRPCMaxRequestsBytes
+	}
+
 	if c.Profiling.Port == 0 {
 		c.Profiling.Port = DefaultProfilingPort
 	}
diff --git a/yorkie/config.sample.yml b/yorkie/config.sample.yml
index f0835fa27..3b19eac55 100644
--- a/yorkie/config.sample.yml
+++ b/yorkie/config.sample.yml
@@ -3,6 +3,9 @@ RPC:
   # Port to listen on for RPC connections (default: 11101).
   Port: 11101
 
+  # MaxRequestBytes is the maximum client request size in bytes the server will accept (default: 4194304, 4MiB).
+  MaxRequestBytes: 4194304
+
   # CertFile is the file containing the TLS certificate.
   CertFile: ""
 
@@ -21,10 +24,10 @@ Profiling:
 Backend:
   # SnapshotThreshold is the threshold that determines if changes should be
   # sent with snapshot when the number of changes is greater than this value.
-  SnapshotThreshold: 500
+  SnapshotThreshold: 1000
 
-  # SnapshotInterval is the interval of changes to create a snapshot.
-  SnapshotInterval: 100
+  # SnapshotInterval is the number of changes to create a snapshot.
+  SnapshotInterval: 5000
 
   # AuthWebhookURL is the URL to send authorization requests to.
   AuthWebhookURL: ""
diff --git a/yorkie/rpc/config.go b/yorkie/rpc/config.go
index 65bd0d6e0..349a8e799 100644
--- a/yorkie/rpc/config.go
+++ b/yorkie/rpc/config.go
@@ -33,9 +33,17 @@ var (
 
 // Config is the configuration for creating a Server instance.
 type Config struct {
-	Port     int
+	// Port is the port number for the RPC server.
+	Port int
+
+	// CertFile is the path to the certificate file.
 	CertFile string
-	KeyFile  string
+
+	// KeyFile is the path to the key file.
+	KeyFile string
+
+	// MaxRequestBytes is the maximum client request size in bytes the server will accept.
+	MaxRequestBytes uint64
 }
 
 // Validate validates the port number and the files for certification.
diff --git a/yorkie/rpc/server.go b/yorkie/rpc/server.go
index 4a3516688..bc15275fd 100644
--- a/yorkie/rpc/server.go
+++ b/yorkie/rpc/server.go
@@ -68,6 +68,8 @@ func NewServer(conf *Config, be *backend.Backend) (*Server, error) {
 		opts = append(opts, grpc.Creds(creds))
 	}
 
+	opts = append(opts, grpc.MaxRecvMsgSize(int(conf.MaxRequestBytes)))
+	opts = append(opts, grpc.MaxSendMsgSize(math.MaxInt32))
 	opts = append(opts, grpc.MaxConcurrentStreams(math.MaxUint32))
 
 	yorkieServiceCtx, yorkieServiceCancel := context.WithCancel(context.Background())
diff --git a/yorkie/rpc/server_test.go b/yorkie/rpc/server_test.go
index 9f0181377..97c7d1d8f 100644
--- a/yorkie/rpc/server_test.go
+++ b/yorkie/rpc/server_test.go
@@ -62,7 +62,8 @@ func TestMain(m *testing.M) {
 	}
 
 	testRPCServer, err = rpc.NewServer(&rpc.Config{
-		Port: helper.RPCPort,
+		Port:            helper.RPCPort,
+		MaxRequestBytes: helper.RPCMaxRequestBytes,
 	}, be)
 	if err != nil {
 		log.Fatal(err)