*: add the http api and new configuration members for ballast object

config/config.go

@@ -184,6 +184,12 @@ type Config struct {
 	// 1. there is a network partition problem between TiDB and PD leader.
 	// 2. there is a network partition problem between TiDB and TiKV leader.
 	EnableForwarding bool `toml:"enable-forwarding" json:"enable-forwarding"`
+	// MaxBallastObjectSize set the max size of the ballast object, the unit is byte.
+	// The default value is the smallest of the following two values: 2GB or
+	// one quarter of the total physical memory in the current system.
+	MaxBallastObjectSize int `toml:"max-ballast-object-size" json:"max-ballast-object-size"`
+	// BallastObjectSize set the initial size of the ballast object, the unit is byte.
+	BallastObjectSize int `toml:"ballast-object-size" json:"ballast-object-size"`
 }
 
 // UpdateTempStoragePath is to update the `TempStoragePath` if port/statusPort was changed

docs/tidb_http_api.md

@@ -522,3 +522,11 @@ timezone.*
     curl -X POST -d "tidb_enable_1pc=0" http://{TiDBIP}:10080/settings
     ```
 
+1. Get/Set the size of the Ballast Object
+
+    ```shell
+    # get current size of the ballast object
+    curl -v http://{TiDBIP}:10080/debug/ballast-object-sz
+    # reset the size of the ballast object (2GB in this example)
+    curl -v -X POST -d "2147483648" http://{TiDBIP}:10080/debug/ballast-object-sz
+    ```

server/http_status.go

@@ -31,6 +31,7 @@ import (
 	rpprof "runtime/pprof"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/gorilla/mux"
@@ -43,6 +44,7 @@ import (
 	"github.com/pingcap/tidb/parser/terror"
 	"github.com/pingcap/tidb/util"
 	"github.com/pingcap/tidb/util/logutil"
+	"github.com/pingcap/tidb/util/memory"
 	"github.com/pingcap/tidb/util/printer"
 	"github.com/pingcap/tidb/util/topsql/tracecpu"
 	"github.com/pingcap/tidb/util/versioninfo"
@@ -194,6 +196,74 @@ func (s *Server) startHTTPServer() {
 	serverMux.HandleFunc("/debug/pprof/profile", tracecpu.ProfileHTTPHandler)
 	serverMux.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
 	serverMux.HandleFunc("/debug/pprof/trace", pprof.Trace)
+	ballast := make([]byte, 0)
+	ballastMaxSz := 1024 * 1024 * 1024 * 2
+	var ballastLock sync.Mutex
+	getBallastSize := func() int {
+		var sz int
+		ballastLock.Lock()
+		sz = len(ballast)
+		ballastLock.Unlock()
+		return sz
+	}
+	setBallastSize := func(newSz int) error {
+		if newSz < 0 {
+			return fmt.Errorf("newSz cannnot be negative: %d", newSz)
+		}
+		if newSz > ballastMaxSz {
+			return fmt.Errorf("newSz cannnot be bigger than %d but it has value %d", ballastMaxSz, newSz)
+		}
+		ballastLock.Lock()
+		ballast = make([]byte, newSz)
+		ballastLock.Unlock()
+		return nil
+	}
+	{
+		if s.cfg.MaxBallastObjectSize > 0 {
+			ballastMaxSz = s.cfg.MaxBallastObjectSize
+		} else {
+			// we try to use the total amount of ram as a reference to set the default ballastMaxSz
+			// since the fatal throw "runtime: out of memory" would never yield to `recover`
+			totalRAMSz, err := memory.MemTotal()
+			if err != nil {
+				logutil.BgLogger().Error("failed to get the total amount of RAM on this system", zap.Error(err))
+			} else {
+				maxSzAdvice := totalRAMSz >> 2
+				if uint64(ballastMaxSz) > maxSzAdvice {
+					ballastMaxSz = int(maxSzAdvice)
+				}
+			}
+		}
+		err := setBallastSize(s.cfg.BallastObjectSize)
+		if err != nil {
+			logutil.BgLogger().Error("set initial ballast object size failed", zap.Error(err))
+		}
+	}
+	serverMux.HandleFunc("/debug/ballast-object-sz", func(w http.ResponseWriter, r *http.Request) {
+		switch r.Method {
+		case http.MethodGet:
+			_, err := w.Write([]byte(strconv.Itoa(getBallastSize())))
+			terror.Log(err)
+		case http.MethodPost:
+			body, err := io.ReadAll(r.Body)
+			if err != nil {
+				terror.Log(err)
+				return
+			}
+			newSz, err := strconv.Atoi(string(body))
+			if err == nil {
+				err = setBallastSize(newSz)
+			}
+			if err != nil {
+				w.WriteHeader(http.StatusBadRequest)
+				errStr := err.Error()
+				if _, err := w.Write([]byte(errStr)); err != nil {
+					terror.Log(err)
+				}
+				return
+			}
+		}
+	})
 	serverMux.HandleFunc("/debug/gogc", func(w http.ResponseWriter, r *http.Request) {
 		switch r.Method {
 		case http.MethodGet: