From 9ee2bb24eda921126d72d2abb4bd1686ff929384 Mon Sep 17 00:00:00 2001
From: Tonis Tiigi <tonistiigi@gmail.com>
Date: Mon, 9 Oct 2023 15:11:22 -0700
Subject: [PATCH] uploadprovider: avoid too big grpc messages on tar upload

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>
---
 session/upload/uploadprovider/provider.go | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/session/upload/uploadprovider/provider.go b/session/upload/uploadprovider/provider.go
index fe7b26a202ba2..d401faa1204f5 100644
--- a/session/upload/uploadprovider/provider.go
+++ b/session/upload/uploadprovider/provider.go
@@ -59,6 +59,23 @@ type writer struct {
 }
 
 func (w *writer) Write(dt []byte) (int, error) {
+	// avoid sending too big messages on grpc stream
+	const maxChunkSize = 3 * 1024 * 1024
+	if len(dt) > maxChunkSize {
+		var n int
+		var err error
+		n, err = w.Write(dt[:maxChunkSize])
+		if err != nil {
+			return n, err
+		}
+		dt = dt[maxChunkSize:]
+		if n2, err := w.Write(dt); err != nil {
+			return n + n2, err
+		} else {
+			n += n2
+		}
+		return n, nil
+	}
 	if err := w.SendMsg(&upload.BytesMessage{Data: dt}); err != nil {
 		return 0, err
 	}