add "setmemlimit" command for utilizing new memory limit feature in Go 1.19+

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
 }

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
+}

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")
+}

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)
 }

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)
 )