From c8c413fa57246b64e9aa0dc2aacbc2a5fa70bcb5 Mon Sep 17 00:00:00 2001
From: Qishuai Liu <lqs@lqs.me>
Date: Fri, 7 Jul 2023 13:18:14 +0900
Subject: [PATCH 1/2] add "setmemlimit" command for utilizing new memory limit
 feature in Go 1.19+

---
 agent/agent.go               | 11 +++++++++++
 agent/memory_limit.go        | 14 ++++++++++++++
 agent/memory_limit_lt1.19.go | 20 ++++++++++++++++++++
 internal/cmd/shared.go       | 27 +++++++++++++++++++++++++++
 signal/signal.go             |  3 +++
 5 files changed, 75 insertions(+)
 create mode 100644 agent/memory_limit.go
 create mode 100644 agent/memory_limit_lt1.19.go

diff --git a/agent/agent.go b/agent/agent.go
index b7978069..d2e5f3fe 100644
--- a/agent/agent.go
+++ b/agent/agent.go
@@ -279,6 +279,17 @@ func handle(conn io.ReadWriter, msg []byte) error {
 			return err
 		}
 		fmt.Fprintf(conn, "New GC percent set to %v. Previous value was %v.\n", perc, debug.SetGCPercent(int(perc)))
+	case signal.SetMemLimit:
+		limit, err := binary.ReadVarint(bufio.NewReader(conn))
+		if err != nil {
+			return err
+		}
+		previous, err := setMemoryLimit(limit)
+		if err != nil {
+			return err
+		}
+		fmt.Fprintf(conn, "New memory limit set to %v. Previous value was %v.\n",
+			formatBytes(uint64(limit)), formatBytes(uint64(previous)))
 	}
 	return nil
 }
diff --git a/agent/memory_limit.go b/agent/memory_limit.go
new file mode 100644
index 00000000..a763c151
--- /dev/null
+++ b/agent/memory_limit.go
@@ -0,0 +1,14 @@
+// Copyright 2023 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+//go:build go1.19
+// +build go1.19
+
+package agent
+
+import "runtime/debug"
+
+func setMemoryLimit(limit int64) (int64, error) {
+	return debug.SetMemoryLimit(limit), nil
+}
diff --git a/agent/memory_limit_lt1.19.go b/agent/memory_limit_lt1.19.go
new file mode 100644
index 00000000..516c7518
--- /dev/null
+++ b/agent/memory_limit_lt1.19.go
@@ -0,0 +1,20 @@
+// Copyright 2023 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+//go:build !go1.19
+// +build !go1.19
+
+package agent
+
+import (
+	"errors"
+	"math"
+)
+
+func setMemoryLimit(limit int64) (int64, error) {
+	if limit < 0 {
+		return math.MaxInt64
+	}
+	return 0, errors.New("memory limit not supported on Go versions before 1.19")
+}
diff --git a/internal/cmd/shared.go b/internal/cmd/shared.go
index fc11e6b2..cea985bf 100644
--- a/internal/cmd/shared.go
+++ b/internal/cmd/shared.go
@@ -9,6 +9,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"math"
 	"net"
 	"os"
 	"os/exec"
@@ -45,6 +46,11 @@ func AgentCommands() []*cobra.Command {
 			short: "Sets the garbage collection target percentage. To completely stop GC, set to 'off'",
 			fn:    setGC,
 		},
+		{
+			name:  "setmemlimit",
+			short: "Sets the memory limit for the process. To disable the limit, set to 'off'",
+			fn:    setMemoryLimit,
+		},
 		{
 			name:  "memstats",
 			short: "Prints the allocation and garbage collection stats.",
@@ -143,6 +149,27 @@ func setGC(addr net.TCPAddr, params []string) error {
 	return cmdWithPrint(addr, signal.SetGCPercent, buf...)
 }
 
+func setMemoryLimit(addr net.TCPAddr, params []string) error {
+	if len(params) != 1 {
+		return errors.New("missing memory limit")
+	}
+	var (
+		limit int64
+		err   error
+	)
+	if strings.ToLower(params[0]) == "off" {
+		limit = math.MaxInt64
+	} else {
+		limit, err = strconv.ParseInt(params[0], 10, strconv.IntSize)
+		if err != nil {
+			return err
+		}
+	}
+	buf := make([]byte, binary.MaxVarintLen64)
+	binary.PutVarint(buf, limit)
+	return cmdWithPrint(addr, signal.SetMemLimit, buf...)
+}
+
 func stackTrace(addr net.TCPAddr, _ []string) error {
 	return cmdWithPrint(addr, signal.StackTrace)
 }
diff --git a/signal/signal.go b/signal/signal.go
index c70764a0..cea69c5c 100644
--- a/signal/signal.go
+++ b/signal/signal.go
@@ -35,4 +35,7 @@ const (
 
 	// SetGCPercent sets the garbage collection target percentage.
 	SetGCPercent = byte(0x10)
+
+	// SetMemLimit sets the memory limit. (Go 1.19+)
+	SetMemLimit = byte(0x11)
 )

From 25b77ae0c9ad2922bf45b22146b052f98274c34d Mon Sep 17 00:00:00 2001
From: Qishuai Liu <lqs@lqs.me>
Date: Fri, 7 Jul 2023 17:44:36 +0900
Subject: [PATCH 2/2] fix return value in memory_limit_lt1.19.go

---
 agent/memory_limit_lt1.19.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/agent/memory_limit_lt1.19.go b/agent/memory_limit_lt1.19.go
index 516c7518..7d6c5f75 100644
--- a/agent/memory_limit_lt1.19.go
+++ b/agent/memory_limit_lt1.19.go
@@ -14,7 +14,7 @@ import (
 
 func setMemoryLimit(limit int64) (int64, error) {
 	if limit < 0 {
-		return math.MaxInt64
+		return math.MaxInt64, nil
 	}
 	return 0, errors.New("memory limit not supported on Go versions before 1.19")
 }