chore(memlimit): Add panic-guard to SetGoMemLimitWithEnv

KimMachineGun · Mar 10, 2023 · 39585f3 · 39585f3
1 parent ac52265
commit 39585f3
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/memlimit/memlimit.go b/memlimit/memlimit.go
@@ -36,6 +36,15 @@ var (
 // If AUTOMEMLIMIT is not set, it defaults to 0.9. (10% is the headroom for memory sources the Go runtime is unaware of.)
 // If GOMEMLIMIT is already set or AUTOMEMLIMIT=off, this function does nothing.
 func SetGoMemLimitWithEnv() {
+	snapshot := debug.SetMemoryLimit(-1)
+	defer func() {
+		err := recover()
+		if err != nil {
+			logger.Printf("panic during SetGoMemLimitWithEnv, rolling back to previous value %d: %v\n", snapshot, err)
+			debug.SetMemoryLimit(snapshot)
+		}
+	}()
+
 	if os.Getenv(envAUTOMEMLIMIT_DEBUG) == "true" {
 		logger = log.Default()
 	}