feat: add x/runtime AutoGOMAXPROCs, PrintStackTrace, SetCurrentUser; add singleflight

Author: sendya

x/runtime/maxprocs.go

@@ -0,0 +1,24 @@
+package runtime
+
+import (
+	"os"
+	"runtime"
+	"runtime/debug"
+	"strconv"
+)
+
+const MaxMemory = 512 << 20 // 512MB
+
+func AutoGOMAXPROCS() (int, int) {
+	// set the max proc = preCore / 2
+	maxThreads := 1000
+	if envStr, ok := os.LookupEnv("APP_MAXTHREADS"); ok {
+		if v, err := strconv.Atoi(envStr); err == nil {
+			maxThreads = v
+		}
+	}
+
+	debug.SetMaxThreads(maxThreads)
+	debug.SetMemoryLimit(MaxMemory)
+	return runtime.GOMAXPROCS(runtime.NumCPU()), maxThreads
+}

x/runtime/maxprocs_test.go

@@ -0,0 +1,46 @@
+package runtime_test
+
+import (
+	"os"
+	"testing"
+
+	"github.com/omalloc/contrib/x/runtime"
+)
+
+func TestAutoGOMAXPROCS(t *testing.T) {
+	tests := []struct {
+		name           string
+		envMaxThreads  string
+		wantMaxThreads int
+	}{
+		{
+			name:           "default_value",
+			envMaxThreads:  "",
+			wantMaxThreads: 1000,
+		},
+		{
+			name:           "custom_max_threads",
+			envMaxThreads:  "500",
+			wantMaxThreads: 500,
+		},
+		{
+			name:           "invalid_env_value",
+			envMaxThreads:  "invalid",
+			wantMaxThreads: 1000,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if tt.envMaxThreads != "" {
+				os.Setenv("APP_MAXTHREADS", tt.envMaxThreads)
+				defer os.Unsetenv("APP_MAXTHREADS")
+			}
+
+			_, gotMaxThreads := runtime.AutoGOMAXPROCS()
+			if gotMaxThreads != tt.wantMaxThreads {
+				t.Errorf("AutoGOMAXPROCS() maxThreads = %v, want %v", gotMaxThreads, tt.wantMaxThreads)
+			}
+		})
+	}
+}

x/runtime/recovered.go

@@ -0,0 +1,30 @@
+package runtime
+
+import (
+	"fmt"
+	"runtime"
+	"strings"
+)
+
+func PrintStackTrace(skip int) string {
+	// Capture the stack trace
+	pc := make([]uintptr, 10)
+	n := runtime.Callers(skip, pc)
+	frames := runtime.CallersFrames(pc[:n])
+
+	// Iterate over the frames and print them
+	sb := strings.Builder{}
+	for {
+		frame, more := frames.Next()
+		if strings.HasPrefix(frame.Function, "runtime.panic") || strings.HasPrefix(frame.Function, "runtime.gopanic") {
+			sb.WriteString("panic: ")
+			continue
+		}
+		sb.WriteString(fmt.Sprintf("%s\n\t%s:%d\n", frame.Function, frame.File, frame.Line))
+		if !more {
+			break
+		}
+	}
+
+	return sb.String()
+}

x/runtime/recovered_test.go

@@ -0,0 +1,60 @@
+package runtime_test
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/omalloc/contrib/x/runtime"
+)
+
+func TestPrintStackTrace(t *testing.T) {
+	// 测试基本堆栈跟踪输出
+	trace := runtime.PrintStackTrace(1)
+
+	// 验证输出包含当前测试函数名
+	if !strings.Contains(trace, "TestPrintStackTrace") {
+		t.Errorf("堆栈跟踪中应该包含测试函数名 'TestPrintStackTrace'，但实际输出为:\n%s", trace)
+	}
+
+	// 验证输出包含文件名和行号
+	if !strings.Contains(trace, "recovered_test.go") {
+		t.Errorf("堆栈跟踪中应该包含测试文件名 'recovered_test.go'，但实际输出为:\n%s", trace)
+	}
+
+	// 验证输出格式
+	lines := strings.Split(strings.TrimSpace(trace), "\n")
+	if len(lines) < 2 {
+		t.Errorf("堆栈跟踪应至少包含两行，但实际输出为:\n%s", trace)
+	}
+
+	// 测试不同的 skip 值
+	trace2 := runtime.PrintStackTrace(2)
+	if trace == trace2 {
+		t.Error("不同的 skip 值应产生不同的堆栈跟踪")
+	}
+}
+
+// 用于测试 panic 情况的辅助函数
+func causePanic() string {
+	panic("测试 panic")
+}
+
+func TestPrintStackTraceWithPanic(t *testing.T) {
+	defer func() {
+		if r := recover(); r != nil {
+			trace := runtime.PrintStackTrace(1)
+
+			// 验证输出包含 panic 前缀
+			if !strings.Contains(trace, "panic:") {
+				t.Errorf("panic 堆栈跟踪应包含 'panic:' 前缀，但实际输出为:\n%s", trace)
+			}
+
+			// 验证输出包含触发 panic 的函数名
+			if !strings.Contains(trace, "causePanic") {
+				t.Errorf("堆栈跟踪应包含 'causePanic' 函数，但实际输出为:\n%s", trace)
+			}
+		}
+	}()
+
+	causePanic()
+}

x/runtime/user.go

@@ -0,0 +1,46 @@
+package runtime
+
+import (
+	"fmt"
+	"os/user"
+	"strconv"
+	"syscall"
+)
+
+func SetCurrentUser(username string) error {
+	if username != "" {
+		current, err := user.Current()
+		if err != nil {
+			return err
+		}
+
+		if current.Username == username {
+			return nil
+		}
+
+		wantedUser, err := user.Lookup(username)
+		if err != nil {
+			return err
+		}
+
+		uid, err := strconv.Atoi(wantedUser.Uid)
+		if err != nil {
+			return fmt.Errorf("error converting UID [%s] to int: %s", wantedUser.Uid, err)
+		}
+
+		gid, err := strconv.Atoi(wantedUser.Gid)
+		if err != nil {
+			return fmt.Errorf("error converting GID [%s] to int: %s", wantedUser.Gid, err)
+		}
+
+		if err = syscall.Setgid(gid); err != nil {
+			return fmt.Errorf("setting group id: %s", err)
+		}
+
+		if err = syscall.Setuid(uid); err != nil {
+			return fmt.Errorf("setting user id: %s", err)
+		}
+	}
+
+	return nil
+}

x/runtime/user_test.go

@@ -0,0 +1,55 @@
+package runtime_test
+
+import (
+	"os/user"
+	"testing"
+
+	"github.com/omalloc/contrib/x/runtime"
+)
+
+func TestSetCurrentUser(t *testing.T) {
+	tests := []struct {
+		name     string
+		username string
+		wantErr  bool
+	}{
+		{
+			name:     "空用户名应该直接返回nil",
+			username: "",
+			wantErr:  false,
+		},
+		{
+			name:     "当前用户名相同应该返回nil",
+			username: getCurrentUsername(t),
+			wantErr:  false,
+		},
+		{
+			name:     "不存在的用户名应该返回错误",
+			username: "nonexistentuser123456789",
+			wantErr:  true,
+		},
+		{
+			name:     "切换用户为daemon",
+			username: "daemon",
+			wantErr:  true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			err := runtime.SetCurrentUser(tt.username)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("SetCurrentUser() error = %v, wantErr %v", err, tt.wantErr)
+			}
+		})
+	}
+}
+
+// 辅助函数：获取当前用户名
+func getCurrentUsername(t *testing.T) string {
+	current, err := user.Current()
+	if err != nil {
+		t.Fatalf("无法获取当前用户: %v", err)
+	}
+	return current.Username
+}