all: initial implementation

Author: changkun

README.md

@@ -1,2 +1,3 @@
 # mkill
+
 mkill limites the number of threads in a Go program

example/main.go

@@ -0,0 +1,17 @@
+package main
+
+import (
+	"time"
+
+	"github.com/changkun/mkill"
+)
+
+func main() {
+	mkill.GOMAXTHREADS(10)
+	for {
+		time.Sleep(time.Second)
+		go func() {
+			time.Sleep(time.Second * 10)
+		}()
+	}
+}

go.mod

@@ -0,0 +1,3 @@
+module github.com/changkun/mkill
+
+go 1.13

mkill.go

@@ -0,0 +1,78 @@
+package mkill
+
+import (
+	"fmt"
+	"os"
+	"os/exec"
+	"runtime"
+	"strconv"
+	"strings"
+	"sync/atomic"
+	"time"
+)
+
+var (
+	pid       = os.Getpid()
+	maxThread = int32(runtime.NumCPU())
+	interval  = time.Second
+	debug     = true
+)
+
+func checkwork() {
+	_, err := getThreads()
+	if err != nil {
+		panic(fmt.Sprintf("mkill: failed to use the library: %v", err))
+	}
+}
+
+func init() {
+	checkwork()
+
+	if debug {
+		fmt.Printf("mkill: pid %v, maxThread %v, interval %v\n", pid, maxThread, interval)
+	}
+	go func() {
+		t := time.NewTicker(interval)
+		for {
+			select {
+			case <-t.C:
+				n, _ := getThreads()
+				nkill := int32(n) - atomic.LoadInt32(&maxThread)
+				if nkill <= 0 {
+					if debug {
+						fmt.Printf("mkill: checked #threads total %v / max %v\n", n, maxThread)
+					}
+					continue
+				}
+				for i := int32(0); i < nkill; i++ {
+					go func() {
+						runtime.LockOSThread()
+					}()
+				}
+				if debug {
+					fmt.Printf("mkill: killing #threads, remaining: %v\n", n)
+				}
+			}
+		}
+	}()
+}
+
+// GOMAXTHREADS change the limites of the maximum threads in runtime
+// and returns the previous number of threads limit
+func GOMAXTHREADS(n int) int {
+	return int(atomic.SwapInt32(&maxThread, int32(n)))
+}
+
+// getThreads returns the number of running threads
+// Linux:
+func getThreads() (int, error) {
+	out, err := exec.Command("bash", "-c", cmdThreads).Output()
+	if err != nil {
+		return 0, fmt.Errorf("mkill: failed to fetch #threads: %v", err)
+	}
+	n, err := strconv.Atoi(strings.TrimSpace(string(out)))
+	if err != nil {
+		return 0, fmt.Errorf("mkill: failed to parse #threads: %v", err)
+	}
+	return n, nil
+}

platform_darwin.go

@@ -0,0 +1,7 @@
+// +build darwin
+
+package mkill
+
+import "fmt"
+
+var cmdThreads = fmt.Sprintf("ps M %d | wc -l", pid)

platform_linux.go

@@ -0,0 +1,7 @@
+// +build linux
+
+package mkill
+
+import "fmt"
+
+var cmdThreads = fmt.Sprintf("ps hH p %d | wc -l", pid)