all: rework entire package

Author: changkun

.github/workflows/mkill.yml

@@ -0,0 +1,37 @@
+name: mkill
+
+on:
+  push:
+    branches: [ master ]
+  pull_request:
+    branches: [ master ]
+
+jobs:
+
+  build:
+    name: Build
+    runs-on: ubuntu-latest
+    steps:
+
+    - name: Set up Go 1.x
+      uses: actions/setup-go@v2
+      with:
+        go-version: ^1.13
+      id: go
+
+    - name: Check out code into the Go module directory
+      uses: actions/checkout@v2
+
+    - name: Get dependencies
+      run: |
+        go get -v -t -d ./...
+
+    - name: Test
+      run: |
+        go test -v -coverprofile=coverage.txt -covermode=atomic ./...
+
+    - name: Upload coverage profile
+      uses: codecov/codecov-action@v1
+      with:
+        token: ${{secrets.CODECOV_TOKEN}}
+        file: coverage.txt

LICENSE

@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2019 Ou Changkun
+Copyright (c) 2020 The golang.design Authors
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

README.md

@@ -1,13 +1,20 @@
 # mkill
 
-mkill limits the number of threads in a Go program
+[![PkgGoDev](https://pkg.go.dev/badge/golang.design/x/mkill)](https://pkg.go.dev/golang.design/x/mkill) [![Go Report Card](https://goreportcard.com/badge/golang.design/x/mkill)](https://goreportcard.com/report/golang.design/x/mkill)
+![mkill](https://github.com/golang-design/mkill/workflows/mkill/badge.svg?branch=master)
 
-## Usage
+Package mkill limits the number of threads in a Go program, without crash the whole program.
 
 ```
-mkill.GOMAXTHREADS(50)
+import "golang.design/x/mkill"
+```
+
+## Quick Start
+
+```
+mkill.GOMAXTHREADS(10)
 ```
 
 ## License
 
-MIT © Changkun Ou
+MIT © The [golang.design](https://golang.design) Authors

example/main.go

@@ -1,3 +1,3 @@
-module github.com/changkun/mkill
+module golang.design/x/mkill
 
 go 1.13

go.mod

@@ -1,78 +1,116 @@
-package mkill
+// Copyright 2020 The golang.design Initiative authors.
+// All rights reserved. Use of this source code is governed
+// by a MIT license that can be found in the LICENSE file.
+
+// Package mkill limits the number of threads in a Go program, without crash the whole program.
+package mkill // import "golang.design/x/mkill"
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"os/exec"
 	"runtime"
 	"strconv"
 	"strings"
+	"sync"
 	"sync/atomic"
 	"time"
 )
 
 var (
 	pid       = os.Getpid()
-	maxThread = int32(runtime.NumCPU())
+	maxThread = int32(runtime.NumCPU()) + 2 // 2 meaning runtime sysmon thread + template thread
 	interval  = time.Second
 	debug     = false
 )
 
+// NumM returns the number of running threads.
+func NumM() int {
+	out, err := exec.Command("bash", "-c", cmdThreads).Output()
+	if err != nil && debug {
+		fmt.Printf("mkill: failed to fetch #threads: %v\n", err)
+		return 0
+	}
+	n, err := strconv.Atoi(strings.TrimSpace(string(out)))
+	if err != nil && debug {
+		fmt.Printf("mkill: failed to parse #threads: %v\n", err)
+		return 0
+	}
+	return n
+}
+
+// GOMAXTHREADS sets the maximum number of system threads that allowed in a Go program
+// and returns the previous setting. If n < 1, it does not change the current setting.
+// The default allowed number of threads of a program is runtime.NumCPU() + 2.
+func GOMAXTHREADS(n int) int {
+	if n < 1 {
+		return int(atomic.LoadInt32(&maxThread))
+	}
+
+	return int(atomic.SwapInt32(&maxThread, int32(n)))
+}
+
+// Wait waits until the number of threads meet the GOMAXTHREADS settings.
+// The function always returns true if the ctx is not canceled.
+// Otherwise returns true only if the Wait is successed in the last check.
+func Wait(ctx context.Context) (ok bool) {
+	for {
+		select {
+		case <-ctx.Done():
+			if NumM() <= GOMAXTHREADS(0) {
+				ok = true
+			}
+			return
+		default:
+			if NumM() > GOMAXTHREADS(0) {
+				continue
+			}
+			ok = true
+			return
+		}
+	}
+}
+
 func checkwork() {
-	_, err := getThreads()
+	_, err := exec.Command("bash", "-c", cmdThreads).Output()
 	if err != nil {
-		panic(fmt.Sprintf("mkill: failed to use the library: %v", err))
+		panic(fmt.Sprintf("mkill: failed to use the package: %v", err))
 	}
 }
 
 func init() {
 	checkwork()
-
 	if debug {
 		fmt.Printf("mkill: pid %v, maxThread %v, interval %v\n", pid, maxThread, interval)
 	}
+
+	wg := sync.WaitGroup{}
 	go func() {
 		t := time.NewTicker(interval)
 		for {
 			select {
 			case <-t.C:
-				n, _ := getThreads()
+				n := NumM()
 				nkill := int32(n) - atomic.LoadInt32(&maxThread)
 				if nkill <= 0 {
 					if debug {
 						fmt.Printf("mkill: checked #threads total %v / max %v\n", n, maxThread)
 					}
 					continue
 				}
+				wg.Add(int(nkill))
 				for i := int32(0); i < nkill; i++ {
 					go func() {
 						runtime.LockOSThread()
+						wg.Done()
 					}()
 				}
+				wg.Wait()
 				if debug {
 					fmt.Printf("mkill: killing #threads, remaining: %v\n", n)
 				}
 			}
 		}
 	}()
 }
-
-// GOMAXTHREADS change the limits 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
-}

mkill.go

@@ -0,0 +1,55 @@
+// Copyright 2020 The golang.design Initiative authors.
+// All rights reserved. Use of this source code is governed
+// by a MIT license that can be found in the LICENSE file.
+
+package mkill_test
+
+import (
+	"context"
+	"fmt"
+	"testing"
+	"time"
+
+	"golang.design/x/mkill"
+)
+
+func TestMKill(t *testing.T) {
+	mkill.GOMAXTHREADS(10)
+
+	// create a lot of threads by sleep gs
+	for i := 0; i < 100000; i++ {
+		go func() {
+			time.Sleep(time.Second * 10)
+		}()
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*100)
+	defer cancel()
+	ok := mkill.Wait(ctx)
+	if !ok {
+		t.Fatal("mkill failed in 100s")
+	}
+}
+
+func ExampleGOMAXTHREADS() {
+	mkill.GOMAXTHREADS(10)
+	// Output:
+}
+
+func ExampleNumM() {
+	mkill.GOMAXTHREADS(10)
+	mkill.Wait(context.Background())
+	fmt.Println(mkill.NumM() <= 10)
+	// Output:
+	// true
+}
+
+func ExampleWait() {
+	mkill.GOMAXTHREADS(10)
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*100)
+	defer cancel()
+	fmt.Println(mkill.Wait(ctx))
+	// Output:
+	// true
+}

mkill_test.go

@@ -1,3 +1,7 @@
+// Copyright 2020 The golang.design Initiative authors.
+// All rights reserved. Use of this source code is governed
+// by a MIT license that can be found in the LICENSE file.
+
 // +build darwin
 
 package mkill

platform_darwin.go

@@ -1,3 +1,7 @@
+// Copyright 2020 The golang.design Initiative authors.
+// All rights reserved. Use of this source code is governed
+// by a MIT license that can be found in the LICENSE file.
+
 // +build linux
 
 package mkill