From 151d837b92cee2c7db226f8dc7d53eea05119111 Mon Sep 17 00:00:00 2001 From: Changkun Ou Date: Fri, 25 Sep 2020 18:29:02 +0200 Subject: [PATCH] all: rework entire package --- .github/workflows/mkill.yml | 37 +++++++++++++++ LICENSE | 2 +- README.md | 15 +++++-- example/main.go | 17 ------- go.mod | 2 +- mkill.go | 90 ++++++++++++++++++++++++++----------- mkill_test.go | 55 +++++++++++++++++++++++ platform_darwin.go | 4 ++ platform_linux.go | 4 ++ 9 files changed, 177 insertions(+), 49 deletions(-) create mode 100644 .github/workflows/mkill.yml delete mode 100644 example/main.go create mode 100644 mkill_test.go diff --git a/.github/workflows/mkill.yml b/.github/workflows/mkill.yml new file mode 100644 index 0000000..841ebea --- /dev/null +++ b/.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 \ No newline at end of file diff --git a/LICENSE b/LICENSE index 0e98686..5c84976 100644 --- a/LICENSE +++ b/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 diff --git a/README.md b/README.md index 8411a84..9c9e7b1 100644 --- a/README.md +++ b/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 \ No newline at end of file +MIT © The [golang.design](https://golang.design) Authors \ No newline at end of file diff --git a/example/main.go b/example/main.go deleted file mode 100644 index 27dcde6..0000000 --- a/example/main.go +++ /dev/null @@ -1,17 +0,0 @@ -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) - }() - } -} diff --git a/go.mod b/go.mod index 886a811..8d8dd81 100644 --- a/go.mod +++ b/go.mod @@ -1,3 +1,3 @@ -module github.com/changkun/mkill +module golang.design/x/mkill go 1.13 diff --git a/mkill.go b/mkill.go index 4873bac..ad06336 100644 --- a/mkill.go +++ b/mkill.go @@ -1,42 +1,97 @@ -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 { @@ -44,11 +99,14 @@ func init() { } 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) } @@ -56,23 +114,3 @@ func init() { } }() } - -// 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 -} diff --git a/mkill_test.go b/mkill_test.go new file mode 100644 index 0000000..29c52fa --- /dev/null +++ b/mkill_test.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 +} diff --git a/platform_darwin.go b/platform_darwin.go index 52db5c9..b97086e 100644 --- a/platform_darwin.go +++ b/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 darwin package mkill diff --git a/platform_linux.go b/platform_linux.go index b7fb3cb..b61991f 100644 --- a/platform_linux.go +++ b/platform_linux.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