├── .github ├── dependabot.yml ├── release.yml └── workflows │ ├── ci.yml │ └── tagpr.yml ├── .gitignore ├── .golangci.yml ├── .octocov.yml ├── .tagpr ├── CHANGELOG.md ├── LICENSE ├── Makefile ├── README.md ├── donegroup.go ├── donegroup_test.go ├── example_test.go ├── go.mod └── go.sum /.github/dependabot.yml: -------------------------------------------------------------------------------- 1 | version: 2 2 | updates: 3 | 4 | - package-ecosystem: "github-actions" 5 | directory: "/" 6 | groups: 7 | dependencies: 8 | patterns: 9 | - "*" 10 | schedule: 11 | interval: "weekly" 12 | time: "08:00" 13 | timezone: "Asia/Tokyo" 14 | commit-message: 15 | prefix: "chore" 16 | include: "scope" 17 | open-pull-requests-limit: 10 18 | assignees: 19 | - "k1LoW" 20 | 21 | - package-ecosystem: "gomod" 22 | directory: "/" 23 | groups: 24 | dependencies: 25 | patterns: 26 | - "*" 27 | schedule: 28 | interval: "weekly" 29 | time: "08:00" 30 | timezone: "Asia/Tokyo" 31 | commit-message: 32 | prefix: "chore" 33 | include: "scope" 34 | open-pull-requests-limit: 10 35 | assignees: 36 | - "k1LoW" 37 | -------------------------------------------------------------------------------- /.github/release.yml: -------------------------------------------------------------------------------- 1 | changelog: 2 | exclude: 3 | labels: 4 | - tagpr 5 | categories: 6 | - title: Breaking Changes 🛠 7 | labels: 8 | - breaking-change 9 | - title: New Features 🎉 10 | labels: 11 | - enhancement 12 | - title: Fix bug 🐛 13 | labels: 14 | - bug 15 | - title: Other Changes 16 | labels: 17 | - "*" 18 | -------------------------------------------------------------------------------- /.github/workflows/ci.yml: -------------------------------------------------------------------------------- 1 | name: CI 2 | 3 | on: 4 | push: 5 | branches: 6 | - main 7 | pull_request: 8 | 9 | jobs: 10 | job-test: 11 | name: Test 12 | runs-on: ubuntu-latest 13 | env: 14 | GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} 15 | steps: 16 | - name: Check out source code 17 | uses: actions/checkout@v4 18 | 19 | - name: Set up Go 20 | uses: actions/setup-go@v5 21 | with: 22 | go-version-file: go.mod 23 | cache: true 24 | 25 | - name: Run lint 26 | uses: reviewdog/action-golangci-lint@v2 27 | with: 28 | fail_on_error: true 29 | golangci_lint_flags: --timeout=5m 30 | 31 | - name: Run tests 32 | run: make ci 33 | 34 | - name: Run octocov 35 | uses: k1LoW/octocov-action@v1 36 | -------------------------------------------------------------------------------- /.github/workflows/tagpr.yml: -------------------------------------------------------------------------------- 1 | name: tagpr 2 | on: 3 | push: 4 | branches: 5 | - main 6 | 7 | jobs: 8 | tagpr: 9 | runs-on: ubuntu-latest 10 | env: 11 | GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} 12 | steps: 13 | - name: Check out source code 14 | uses: actions/checkout@v4 15 | 16 | - id: run-tagpr 17 | name: Run tagpr 18 | uses: Songmu/tagpr@v1 19 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | coverage.out 2 | -------------------------------------------------------------------------------- /.golangci.yml: -------------------------------------------------------------------------------- 1 | linters: 2 | fast: false 3 | enable: 4 | - misspell 5 | - gosec 6 | - godot 7 | - revive 8 | linters-settings: 9 | errcheck: 10 | check-type-assertions: true 11 | misspell: 12 | locale: US 13 | ignore-words: [] 14 | revive: 15 | rules: 16 | - name: unexported-return 17 | disabled: true 18 | - name: exported 19 | disabled: false 20 | -------------------------------------------------------------------------------- /.octocov.yml: -------------------------------------------------------------------------------- 1 | # generated by octocov init 2 | coverage: 3 | if: true 4 | codeToTestRatio: 5 | code: 6 | - '**/*.go' 7 | - '!**/*_test.go' 8 | test: 9 | - '**/*_test.go' 10 | testExecutionTime: 11 | if: true 12 | diff: 13 | datastores: 14 | - artifact://${GITHUB_REPOSITORY} 15 | comment: 16 | if: is_pull_request 17 | summary: 18 | if: true 19 | report: 20 | if: is_default_branch 21 | datastores: 22 | - artifact://${GITHUB_REPOSITORY} 23 | -------------------------------------------------------------------------------- /.tagpr: -------------------------------------------------------------------------------- 1 | # config file for the tagpr in git config format 2 | # The tagpr generates the initial configuration, which you can rewrite to suit your environment. 3 | # CONFIGURATIONS: 4 | # tagpr.releaseBranch 5 | # Generally, it is "main." It is the branch for releases. The pcpr tracks this branch, 6 | # creates or updates a pull request as a release candidate, or tags when they are merged. 7 | # 8 | # tagpr.versionFile 9 | # Versioning file containing the semantic version needed to be updated at release. 10 | # It will be synchronized with the "git tag". 11 | # Often this is a meta-information file such as gemspec, setup.cfg, package.json, etc. 12 | # Sometimes the source code file, such as version.go or Bar.pm, is used. 13 | # If you do not want to use versioning files but only git tags, specify the "-" string here. 14 | # You can specify multiple version files by comma separated strings. 15 | # 16 | # tagpr.vPrefix 17 | # Flag whether or not v-prefix is added to semver when git tagging. (e.g. v1.2.3 if true) 18 | # This is only a tagging convention, not how it is described in the version file. 19 | # 20 | # tagpr.changelog (Optional) 21 | # Flag whether or not changelog is added or changed during the release. 22 | # 23 | # tagpr.command (Optional) 24 | # Command to change files just before release. 25 | # 26 | # tagpr.tmplate (Optional) 27 | # Pull request template in go template format 28 | # 29 | # tagpr.release (Optional) 30 | # GitHub Release creation behavior after tagging [true, draft, false] 31 | # If this value is not set, the release is to be created. 32 | [tagpr] 33 | vPrefix = true 34 | releaseBranch = main 35 | command = "make prerelease_for_tagpr" 36 | versionFile = - 37 | -------------------------------------------------------------------------------- /CHANGELOG.md: -------------------------------------------------------------------------------- 1 | # Changelog 2 | 3 | ## [v1.10.2](https://github.com/k1LoW/donegroup/compare/v1.10.1...v1.10.2) - 2024-06-08 4 | ### Other Changes 5 | - Use context.AfterFunc instead of go func() by @k1LoW in https://github.com/k1LoW/donegroup/pull/52 6 | 7 | ## [v1.10.1](https://github.com/k1LoW/donegroup/compare/v1.10.0...v1.10.1) - 2024-06-08 8 | ### Other Changes 9 | - Refactor withDoneGroup by @k1LoW in https://github.com/k1LoW/donegroup/pull/51 10 | 11 | ## [v1.10.0](https://github.com/k1LoW/donegroup/compare/v1.9.0...v1.10.0) - 2024-06-05 12 | ### New Features 🎉 13 | - Add `WithoutCancel` by @k1LoW in https://github.com/k1LoW/donegroup/pull/49 14 | 15 | ## [v1.9.0](https://github.com/k1LoW/donegroup/compare/v1.8.1...v1.9.0) - 2024-06-04 16 | ### Breaking Changes 🛠 17 | - Change the behaviour of `donegroup.Cancel` significantly. by @k1LoW in https://github.com/k1LoW/donegroup/pull/47 18 | 19 | ## [v1.8.1](https://github.com/k1LoW/donegroup/compare/v1.8.0...v1.8.1) - 2024-06-02 20 | ### Fix bug 🐛 21 | - Cleanup functions should be executed immediately when the context is done. by @k1LoW in https://github.com/k1LoW/donegroup/pull/44 22 | ### Other Changes 23 | - doneGroup.ctxw is not used anymore, so remove it. by @k1LoW in https://github.com/k1LoW/donegroup/pull/45 24 | 25 | ## [v1.8.0](https://github.com/k1LoW/donegroup/compare/v1.7.0...v1.8.0) - 2024-06-02 26 | ### Breaking Changes 🛠 27 | - If timeout is reached, it should not be waited for. by @k1LoW in https://github.com/k1LoW/donegroup/pull/39 28 | - Functions registered in Cleanup no longer need to do context handling. by @k1LoW in https://github.com/k1LoW/donegroup/pull/42 29 | ### Other Changes 30 | - Use context.WithoutCancel instead of context.Background in the donegroup package. by @k1LoW in https://github.com/k1LoW/donegroup/pull/41 31 | 32 | ## [v1.7.0](https://github.com/k1LoW/donegroup/compare/v1.6.0...v1.7.0) - 2024-06-02 33 | ### New Features 🎉 34 | - Add `CancelWith*Cause` by @k1LoW in https://github.com/k1LoW/donegroup/pull/37 35 | ### Fix bug 🐛 36 | - Fix doneGroup._ctx tree by @k1LoW in https://github.com/k1LoW/donegroup/pull/35 37 | ### Other Changes 38 | - Use sync.WaitGroup instead of errgroup.Group by @k1LoW in https://github.com/k1LoW/donegroup/pull/32 39 | - Remove unnecessary for loop by @k1LoW in https://github.com/k1LoW/donegroup/pull/33 40 | - Use context.CancelCauseFunc by @k1LoW in https://github.com/k1LoW/donegroup/pull/34 41 | - Fix cancel timing by @k1LoW in https://github.com/k1LoW/donegroup/pull/36 42 | 43 | ## [v1.6.0](https://github.com/k1LoW/donegroup/compare/v1.5.1...v1.6.0) - 2024-05-21 44 | ### New Features 🎉 45 | - Add WithCancelCause by @k1LoW in https://github.com/k1LoW/donegroup/pull/27 46 | - Add `WithDeadline` and `WithTimeout` by @k1LoW in https://github.com/k1LoW/donegroup/pull/30 47 | ### Other Changes 48 | - chore(deps): bump golang.org/x/sync from 0.6.0 to 0.7.0 in the dependencies group by @dependabot in https://github.com/k1LoW/donegroup/pull/29 49 | - chore(deps): bump actions/setup-go from 4 to 5 in the dependencies group by @dependabot in https://github.com/k1LoW/donegroup/pull/28 50 | 51 | ## [v1.5.1](https://github.com/k1LoW/donegroup/compare/v1.5.0...v1.5.1) - 2024-04-04 52 | 53 | ## [v1.5.0](https://github.com/k1LoW/donegroup/compare/v1.4.0...v1.5.0) - 2024-04-04 54 | ### New Features 🎉 55 | - Add `donegroup.Go` by @k1LoW in https://github.com/k1LoW/donegroup/pull/23 56 | 57 | ## [v1.4.0](https://github.com/k1LoW/donegroup/compare/v1.3.0...v1.4.0) - 2024-02-07 58 | ### Breaking Changes 🛠 59 | - Always execute all cleanup functions. by @k1LoW in https://github.com/k1LoW/donegroup/pull/21 60 | 61 | ## [v1.3.0](https://github.com/k1LoW/donegroup/compare/v1.2.0...v1.3.0) - 2024-02-07 62 | ### Breaking Changes 🛠 63 | - Add Awaitable by @k1LoW in https://github.com/k1LoW/donegroup/pull/19 64 | ### Other Changes 65 | - Add ErrNotContainDoneGroup by @k1LoW in https://github.com/k1LoW/donegroup/pull/17 66 | 67 | ## [v1.2.0](https://github.com/k1LoW/donegroup/compare/v1.1.0...v1.2.0) - 2024-02-07 68 | ### New Features 🎉 69 | - Add Cancel for canceling context and waiting for cleanup functions at once. by @k1LoW in https://github.com/k1LoW/donegroup/pull/16 70 | 71 | ## [v1.1.0](https://github.com/k1LoW/donegroup/compare/v1.0.0...v1.1.0) - 2024-02-07 72 | ### New Features 🎉 73 | - Add donegroup.Awaiter by @k1LoW in https://github.com/k1LoW/donegroup/pull/14 74 | ### Other Changes 75 | - Add test for no cleanup functions by @k1LoW in https://github.com/k1LoW/donegroup/pull/12 76 | 77 | ## [v1.0.0](https://github.com/k1LoW/donegroup/compare/v0.2.3...v1.0.0) - 2024-02-06 78 | ### Other Changes 79 | - Add example by @k1LoW in https://github.com/k1LoW/donegroup/pull/11 80 | 81 | ## [v0.2.3](https://github.com/k1LoW/donegroup/compare/v0.2.2...v0.2.3) - 2024-02-05 82 | ### Fix bug 🐛 83 | - Fix typo by @k1LoW in https://github.com/k1LoW/donegroup/pull/9 84 | 85 | ## [v0.2.2](https://github.com/k1LoW/donegroup/compare/v0.2.1...v0.2.2) - 2024-02-05 86 | ### New Features 🎉 87 | - Add WaitWithTimeout* by @k1LoW in https://github.com/k1LoW/donegroup/pull/7 88 | 89 | ## [v0.2.1](https://github.com/k1LoW/donegroup/compare/v0.2.0...v0.2.1) - 2024-02-05 90 | 91 | ## [v0.2.0](https://github.com/k1LoW/donegroup/compare/v0.1.0...v0.2.0) - 2024-02-05 92 | ### Breaking Changes 🛠 93 | - Add WaitWithTimeout by @k1LoW in https://github.com/k1LoW/donegroup/pull/2 94 | - Add WaitWithContext instead of WaitWithTimeout by @k1LoW in https://github.com/k1LoW/donegroup/pull/4 95 | 96 | ## [v0.0.1](https://github.com/k1LoW/donegroup/commits/v0.0.1) - 2024-02-05 97 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | 3 | Copyright © 2024 Ken'ichiro Oyama 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in 13 | all copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 21 | THE SOFTWARE. 22 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | default: test 2 | 3 | ci: test race 4 | 5 | test: 6 | go test ./... -coverprofile=coverage.out -covermode=count -count=1 7 | 8 | race: 9 | go test ./... -race -count=1 -run Test 10 | 11 | lint: 12 | golangci-lint run ./... 13 | 14 | depsdev: 15 | go install github.com/Songmu/ghch/cmd/ghch@latest 16 | go install github.com/Songmu/gocredits/cmd/gocredits@latest 17 | 18 | prerelease: 19 | git pull origin main --tag 20 | go mod tidy 21 | ghch -w -N ${VER} 22 | gocredits . w 23 | git add CHANGELOG.md CREDITS go.mod go.sum 24 | git commit -m'Bump up version number' 25 | git tag ${VER} 26 | 27 | prerelease_for_tagpr: 28 | gocredits . -w 29 | git add CHANGELOG.md CREDITS go.mod go.sum 30 | 31 | .PHONY: default test 32 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # donegroup [![Go Reference](https://pkg.go.dev/badge/github.com/k1LoW/donegroup.svg)](https://pkg.go.dev/github.com/k1LoW/donegroup) [![CI](https://github.com/k1LoW/donegroup/actions/workflows/ci.yml/badge.svg)](https://github.com/k1LoW/donegroup/actions/workflows/ci.yml) ![Coverage](https://raw.githubusercontent.com/k1LoW/octocovs/main/badges/k1LoW/donegroup/coverage.svg) ![Code to Test Ratio](https://raw.githubusercontent.com/k1LoW/octocovs/main/badges/k1LoW/donegroup/ratio.svg) ![Test Execution Time](https://raw.githubusercontent.com/k1LoW/octocovs/main/badges/k1LoW/donegroup/time.svg) 2 | 3 | `donegroup` is a package that provides a graceful cleanup transaction to context.Context when the context is canceled ( **done** ). 4 | 5 | > sync.WaitGroup + <-ctx.Done() = donegroup 6 | 7 | ## Usage 8 | 9 | Use [donegroup.WithCancel](https://pkg.go.dev/github.com/k1LoW/donegroup#WithCancel) instead of [context.WithCancel](https://pkg.go.dev/context#WithCancel). 10 | 11 | Then, it can wait for the cleanup processes associated with the context using [donegroup.Wait](https://pkg.go.dev/github.com/k1LoW/donegroup#Wait). 12 | 13 | ``` go 14 | // before 15 | ctx, cancel := context.WithCancel(context.Background()) 16 | defer cancel() 17 | ``` 18 | 19 | ↓ 20 | 21 | ``` go 22 | // after 23 | ctx, cancel := donegroup.WithCancel(context.Background()) 24 | defer func() { 25 | cancel() 26 | if err := donegroup.Wait(ctx); err != nil { 27 | log.Fatal(err) 28 | } 29 | }() 30 | ``` 31 | 32 | ### [donegroup.Cleanup](https://pkg.go.dev/github.com/k1LoW/donegroup#Cleanup) ( Basic usage ) 33 | 34 | ``` mermaid 35 | gantt 36 | dateFormat mm 37 | axisFormat 38 | 39 | section main 40 | donegroup.WithCancel :milestone, m1, 00, 0m 41 | cancel context using cancel() : milestone, m2, 03, 0m 42 | donegroup.Wait :milestone, m3, 04, 0m 43 | waiting for all registered funcs to finish :b, 04, 2m 44 | 45 | section cleanup func1 46 | register func1 using donegroup.Cleanup :milestone, m1, 01, 0m 47 | waiting for context cancellation :a, 01, 2m 48 | executing func1 :active, b, 03, 3m 49 | 50 | section cleanup func2 51 | register func2 using donegroup.Cleanup :milestone, m3, 02, 0m 52 | waiting for context cancellation :a, 02, 1m 53 | executing func2 :active, b, 03, 2m 54 | ``` 55 | 56 | ```go 57 | package main 58 | 59 | import ( 60 | "context" 61 | "fmt" 62 | "log" 63 | "time" 64 | 65 | "github.com/k1LoW/donegroup" 66 | ) 67 | 68 | func main() { 69 | ctx, cancel := donegroup.WithCancel(context.Background()) 70 | 71 | // Cleanup process func1 of some kind 72 | if err := donegroup.Cleanup(ctx, func() error { 73 | fmt.Println("cleanup func1") 74 | return nil 75 | }); err != nil { 76 | log.Fatal(err) 77 | } 78 | 79 | // Cleanup process func2 of some kind 80 | if err := donegroup.Cleanup(ctx, func() error { 81 | time.Sleep(1 * time.Second) 82 | fmt.Println("cleanup func2") 83 | return nil 84 | }); err != nil { 85 | log.Fatal(err) 86 | } 87 | 88 | defer func() { 89 | cancel() 90 | if err := donegroup.Wait(ctx); err != nil { 91 | log.Fatal(err) 92 | } 93 | }() 94 | 95 | // Main process of some kind 96 | fmt.Println("main") 97 | 98 | // Output: 99 | // main finish 100 | // cleanup func1 101 | // cleanup func2 102 | } 103 | ``` 104 | 105 | [dongroup.Cleanup](https://pkg.go.dev/github.com/k1LoW/donegroup#Cleanup) is similar in usage to [testing.T.Cleanup](https://pkg.go.dev/testing#T.Cleanup), but the order of execution is not guaranteed. 106 | 107 | ### [donegroup.WaitWithTimeout](https://pkg.go.dev/github.com/k1LoW/donegroup#WaitWithTimeout) ( Wait for a specified duration ) 108 | 109 | Using [donegroup.WaitWithTimeout](https://pkg.go.dev/github.com/k1LoW/donegroup#WaitWithTimeout), it is possible to set a timeout for the cleanup processes. 110 | 111 | Note that each cleanup process must handle its own context argument. 112 | 113 | ``` mermaid 114 | gantt 115 | dateFormat mm 116 | axisFormat 117 | 118 | section main 119 | donegroup.WithCancel :milestone, m1, 00, 0m 120 | cancel context using cancel() : milestone, m2, 03, 0m 121 | donegroup.WaitWithTimeout :milestone, m3, 04, 0m 122 | waiting for all registered funcs to finish :b, 04, 1m 123 | timeout :milestone, m4, 05, 0m 124 | 125 | section cleanup func 126 | register func using donegroup.Cleanup :milestone, m1, 01, 0m 127 | waiting for context cancellation :a, 01, 2m 128 | executing func :active, b, 03, 3m 129 | ``` 130 | 131 | ```go 132 | ctx, cancel := donegroup.WithCancel(context.Background()) 133 | 134 | // Cleanup process of some kind 135 | if err := donegroup.Cleanup(ctx, func() error { 136 | fmt.Println("cleanup start") 137 | for i := 0; i < 10; i++ { 138 | time.Sleep(2 * time.Millisecond) 139 | } 140 | fmt.Println("cleanup finish") 141 | return nil 142 | }); err != nil { 143 | log.Fatal(err) 144 | } 145 | 146 | defer func() { 147 | cancel() 148 | timeout := 5 * time.Millisecond 149 | if err := WaitWithTimeout(ctx, timeout); err != nil { 150 | fmt.Println(err) 151 | } 152 | }() 153 | 154 | // Main process of some kind 155 | fmt.Println("main start") 156 | 157 | fmt.Println("main finish") 158 | 159 | // Output: 160 | // main start 161 | // main finish 162 | // cleanup start 163 | // context deadline exceeded 164 | ``` 165 | 166 | ### [donegroup.Awaiter](https://pkg.go.dev/github.com/k1LoW/donegroup#Awaiter) 167 | 168 | In addition to using [donegroup.Cleanup](https://pkg.go.dev/github.com/k1LoW/donegroup#Cleanup) to register a cleanup function after context cancellation, it is possible to use [donegroup.Awaiter](https://pkg.go.dev/github.com/k1LoW/donegroup#Awaiter) to make the execution of an arbitrary process wait. 169 | 170 | ``` mermaid 171 | gantt 172 | dateFormat mm 173 | axisFormat 174 | 175 | section main 176 | donegroup.WithCancel :milestone, m1, 00, 0m 177 | cancel context using cancel() : milestone, m2, 03, 0m 178 | donegroup.Wait :milestone, m3, 04, 0m 179 | waiting for all registered funcs to finish :b, 04, 1m 180 | 181 | section func on goroutine 182 | executing go func :active, b, 01, 4m 183 | donegroup.Awaiter :milestone, m3, 01, 0m 184 | completed() :milestone, m3, 05, 0m 185 | ``` 186 | 187 | ``` go 188 | ctx, cancel := donegroup.WithCancel(context.Background()) 189 | 190 | go func() { 191 | completed, err := donegroup.Awaiter(ctx) 192 | if err != nil { 193 | log.Fatal(err) 194 | return 195 | } 196 | time.Sleep(1000 * time.Millisecond) 197 | fmt.Println("do something") 198 | completed() 199 | }() 200 | 201 | // Main process of some kind 202 | fmt.Println("main") 203 | time.Sleep(10 * time.Millisecond) 204 | 205 | cancel() 206 | if err := donegroup.Wait(ctx); err != nil { 207 | log.Fatal(err) 208 | } 209 | 210 | fmt.Println("finish") 211 | 212 | // Output: 213 | // main 214 | // do something 215 | // finish 216 | ``` 217 | 218 | It is also possible to guarantee the execution of a function block using `defer donegroup.Awaitable(ctx)()`. 219 | 220 | ``` go 221 | go func() { 222 | defer donegroup.Awaitable(ctx)() 223 | time.Sleep(1000 * time.Millisecond) 224 | fmt.Println("do something") 225 | }() 226 | ``` 227 | 228 | ### [donegroup.Go](https://pkg.go.dev/github.com/k1LoW/donegroup#Go) ( Syntax sugar for `go func()` and donegroup.Awaiter ) 229 | 230 | [donegroup.Go](https://pkg.go.dev/github.com/k1LoW/donegroup#Go) can execute arbitrary process asynchronously while still waiting for it to finish, similar to [donegroup.Awaiter](https://pkg.go.dev/github.com/k1LoW/donegroup#Awaiter). 231 | 232 | ``` mermaid 233 | gantt 234 | dateFormat mm 235 | axisFormat 236 | 237 | section main 238 | donegroup.WithCancel :milestone, m1, 00, 0m 239 | cancel context using cancel() : milestone, m2, 03, 0m 240 | donegroup.Wait :milestone, m3, 04, 0m 241 | waiting for all registered funcs to finish :b, 04, 1m 242 | 243 | section func on donegroup.Go 244 | register and execute func using donegroup.Go :milestone, m3, 01, 0m 245 | executing func :active, b, 01, 4m 246 | ``` 247 | 248 | ``` go 249 | donegroup.Go(ctx, func() error { 250 | time.Sleep(1000 * time.Millisecond) 251 | fmt.Println("do something") 252 | return nil 253 | }() 254 | ``` 255 | 256 | Also, with [donegroup.Go](https://pkg.go.dev/github.com/k1LoW/donegroup#Go), the error can be received via [donegroup.Wait](https://pkg.go.dev/github.com/k1LoW/donegroup#Wait). 257 | 258 | ### [donegroup.Cancel](https://pkg.go.dev/github.com/k1LoW/donegroup#Cancel) 259 | 260 | [donegroup.Cancel](https://pkg.go.dev/github.com/k1LoW/donegroup#Cancel) can cancel the context. 261 | 262 | Can be cancelled anywhere with the context. 263 | 264 | ``` go 265 | ctx, _ := donegroup.WithCancel(context.Background()) 266 | 267 | defer func() { 268 | if err := donegroup.Cancel(ctx); err != nil { 269 | log.Fatal(err) 270 | } 271 | if err := donegroup.Wait(ctx); err != nil { 272 | log.Fatal(err) 273 | } 274 | }() 275 | ``` 276 | -------------------------------------------------------------------------------- /donegroup.go: -------------------------------------------------------------------------------- 1 | package donegroup 2 | 3 | import ( 4 | "context" 5 | "errors" 6 | "sync" 7 | "time" 8 | ) 9 | 10 | var doneGroupKey = struct{}{} 11 | var ErrNotContainDoneGroup = errors.New("donegroup: context does not contain a doneGroup. Use donegroup.With* to create a context with a doneGroup") 12 | 13 | // doneGroup is cleanup function groups per Context. 14 | type doneGroup struct { 15 | cancel context.CancelCauseFunc 16 | cleanupGroups []*sync.WaitGroup 17 | errors error 18 | mu sync.Mutex 19 | } 20 | 21 | // WithCancel returns a copy of parent with a new Done channel and a doneGroup. 22 | func WithCancel(ctx context.Context) (context.Context, context.CancelFunc) { 23 | return WithCancelWithKey(ctx, doneGroupKey) 24 | } 25 | 26 | // WithDeadline returns a copy of parent with a new Done channel and a doneGroup. 27 | // If the deadline is exceeded, the cause is set to context.DeadlineExceeded. 28 | func WithDeadline(ctx context.Context, d time.Time) (context.Context, context.CancelFunc) { 29 | return WithDeadlineCause(ctx, d, nil) 30 | } 31 | 32 | // WithTimeout returns a copy of parent with a new Done channel and a doneGroup. 33 | // If the timeout is exceeded, the cause is set to context.DeadlineExceeded. 34 | func WithTimeout(ctx context.Context, timeout time.Duration) (context.Context, context.CancelFunc) { 35 | return WithTimeoutCause(ctx, timeout, nil) 36 | } 37 | 38 | // WithCancelCause returns a copy of parent with a new Done channel and a doneGroup. 39 | func WithCancelCause(ctx context.Context) (context.Context, context.CancelCauseFunc) { 40 | return WithCancelCauseWithKey(ctx, doneGroupKey) 41 | } 42 | 43 | // WithDeadlineCause returns a copy of parent with a new Done channel and a doneGroup. 44 | func WithDeadlineCause(ctx context.Context, d time.Time, cause error) (context.Context, context.CancelFunc) { 45 | return WithDeadlineCauseWithKey(ctx, d, cause, doneGroupKey) 46 | } 47 | 48 | // WithTimeoutCause returns a copy of parent with a new Done channel and a doneGroup. 49 | func WithTimeoutCause(ctx context.Context, timeout time.Duration, cause error) (context.Context, context.CancelFunc) { 50 | return WithTimeoutCauseWithKey(ctx, timeout, cause, doneGroupKey) 51 | } 52 | 53 | // WithoutCancel returns a copy of parent that is not canceled when parent is canceled and does not have a doneGroup. 54 | func WithoutCancel(ctx context.Context) context.Context { 55 | return WithoutCancelWithKey(ctx, doneGroupKey) 56 | } 57 | 58 | // WithoutCancelWithKey returns a copy of parent that is not canceled when parent is canceled and does not have a doneGroup. 59 | func WithoutCancelWithKey(ctx context.Context, key any) context.Context { 60 | return context.WithValue(context.WithoutCancel(ctx), key, nil) 61 | } 62 | 63 | // WithCancelWithKey returns a copy of parent with a new Done channel and a doneGroup. 64 | func WithCancelWithKey(ctx context.Context, key any) (context.Context, context.CancelFunc) { 65 | ctx, cancelCause := WithCancelCauseWithKey(ctx, key) 66 | return ctx, func() { cancelCause(nil) } 67 | } 68 | 69 | // WithDeadlineWithKey returns a copy of parent with a new Done channel and a doneGroup. 70 | func WithDeadlineWithKey(ctx context.Context, d time.Time, key any) (context.Context, context.CancelFunc) { 71 | return WithDeadlineCauseWithKey(ctx, d, nil, key) 72 | } 73 | 74 | // WithTimeoutWithKey returns a copy of parent with a new Done channel and a doneGroup. 75 | func WithTimeoutWithKey(ctx context.Context, timeout time.Duration, key any) (context.Context, context.CancelFunc) { 76 | return WithTimeoutCauseWithKey(ctx, timeout, nil, key) 77 | } 78 | 79 | // WithCancelCauseWithKey returns a copy of parent with a new Done channel and a doneGroup. 80 | func WithCancelCauseWithKey(ctx context.Context, key any) (context.Context, context.CancelCauseFunc) { 81 | ctx, cancelCause := context.WithCancelCause(ctx) 82 | return withDoneGroup(ctx, cancelCause, key), cancelCause 83 | } 84 | 85 | // WithDeadlineCauseWithKey returns a copy of parent with a new Done channel and a doneGroup. 86 | func WithDeadlineCauseWithKey(ctx context.Context, d time.Time, cause error, key any) (context.Context, context.CancelFunc) { 87 | ctx, cancelCause := context.WithCancelCause(ctx) 88 | ctx, cancel := context.WithDeadlineCause(ctx, d, cause) 89 | ctx = withDoneGroup(ctx, cancelCause, key) 90 | return ctx, cancel 91 | } 92 | 93 | // WithTimeoutCauseWithKey returns a copy of parent with a new Done channel and a doneGroup. 94 | func WithTimeoutCauseWithKey(ctx context.Context, timeout time.Duration, cause error, key any) (context.Context, context.CancelFunc) { 95 | return WithDeadlineCauseWithKey(ctx, time.Now().Add(timeout), cause, key) 96 | } 97 | 98 | // Cleanup registers a function to be called when the context is canceled. 99 | func Cleanup(ctx context.Context, f func() error) error { 100 | return CleanupWithKey(ctx, doneGroupKey, f) 101 | } 102 | 103 | // CleanupWithKey Cleanup registers a function to be called when the context is canceled. 104 | func CleanupWithKey(ctx context.Context, key any, f func() error) error { 105 | dg, ok := ctx.Value(key).(*doneGroup) 106 | if !ok { 107 | return ErrNotContainDoneGroup 108 | } 109 | 110 | rootWg := dg.cleanupGroups[0] 111 | dg.mu.Lock() 112 | rootWg.Add(1) 113 | dg.mu.Unlock() 114 | 115 | _ = context.AfterFunc(ctx, func() { 116 | if err := f(); err != nil { 117 | dg.mu.Lock() 118 | dg.errors = errors.Join(dg.errors, err) 119 | dg.mu.Unlock() 120 | } 121 | rootWg.Done() 122 | }) 123 | return nil 124 | } 125 | 126 | // Wait blocks until the context is canceled. Then calls the function registered by Cleanup. 127 | func Wait(ctx context.Context) error { 128 | return WaitWithKey(ctx, doneGroupKey) 129 | } 130 | 131 | // WaitWithTimeout blocks until the context (ctx) is canceled. Then calls the function registered by Cleanup with timeout. 132 | func WaitWithTimeout(ctx context.Context, timeout time.Duration) error { 133 | return WaitWithTimeoutAndKey(ctx, timeout, doneGroupKey) 134 | } 135 | 136 | // WaitWithContext blocks until the context (ctx) is canceled. Then calls the function registered by Cleanup with context (ctxw). 137 | func WaitWithContext(ctx, ctxw context.Context) error { 138 | return WaitWithContextAndKey(ctx, ctxw, doneGroupKey) 139 | } 140 | 141 | // Cancel cancels the context. Then calls the function registered by Cleanup. 142 | func Cancel(ctx context.Context) error { 143 | return CancelWithKey(ctx, doneGroupKey) 144 | } 145 | 146 | // CancelWithCause cancels the context with cause. Then calls the function registered by Cleanup. 147 | func CancelWithCause(ctx context.Context, cause error) error { 148 | return CancelWithCauseAndKey(ctx, cause, doneGroupKey) 149 | } 150 | 151 | // WaitWithKey blocks until the context is canceled. Then calls the function registered by Cleanup. 152 | func WaitWithKey(ctx context.Context, key any) error { 153 | return WaitWithContextAndKey(ctx, context.WithoutCancel(ctx), key) 154 | } 155 | 156 | // WaitWithTimeoutAndKey blocks until the context is canceled. Then calls the function registered by Cleanup with timeout. 157 | func WaitWithTimeoutAndKey(ctx context.Context, timeout time.Duration, key any) error { 158 | ctxw, cancel := context.WithTimeout(context.WithoutCancel(ctx), timeout) 159 | defer cancel() 160 | return WaitWithContextAndKey(ctx, ctxw, key) 161 | } 162 | 163 | // WaitWithContextAndKey blocks until the context is canceled. Then calls the function registered by Cleanup with context (ctxx). 164 | func WaitWithContextAndKey(ctx, ctxw context.Context, key any) error { 165 | dg, ok := ctx.Value(key).(*doneGroup) 166 | if !ok { 167 | return ErrNotContainDoneGroup 168 | } 169 | <-ctx.Done() 170 | wg := &sync.WaitGroup{} 171 | for _, g := range dg.cleanupGroups { 172 | wg.Add(1) 173 | dg.mu.Lock() 174 | go func() { 175 | g.Wait() 176 | wg.Done() 177 | }() 178 | dg.mu.Unlock() 179 | } 180 | ch := make(chan struct{}) 181 | go func() { 182 | wg.Wait() 183 | close(ch) 184 | }() 185 | select { 186 | case <-ch: 187 | case <-ctxw.Done(): 188 | dg.mu.Lock() 189 | defer dg.mu.Unlock() 190 | dg.errors = errors.Join(dg.errors, ctxw.Err()) 191 | } 192 | return dg.errors 193 | } 194 | 195 | // CancelWithKey cancels the context. 196 | func CancelWithKey(ctx context.Context, key any) error { 197 | return CancelWithCauseAndKey(ctx, nil, key) 198 | } 199 | 200 | // CancelWithCauseAndKey cancels the context with cause. 201 | func CancelWithCauseAndKey(ctx context.Context, cause error, key any) error { 202 | dg, ok := ctx.Value(key).(*doneGroup) 203 | if !ok { 204 | return ErrNotContainDoneGroup 205 | } 206 | dg.cancel(cause) 207 | return nil 208 | } 209 | 210 | // Awaiter returns a function that guarantees execution of the process until it is called. 211 | // Note that if the timeout of WaitWithTimeout has passed (or the context of WaitWithContext has canceled), it will not wait. 212 | func Awaiter(ctx context.Context) (completed func(), err error) { 213 | return AwaiterWithKey(ctx, doneGroupKey) 214 | } 215 | 216 | // AwaiterWithKey returns a function that guarantees execution of the process until it is called. 217 | // Note that if the timeout of WaitWithTimeout has passed (or the context of WaitWithContext has canceled), it will not wait. 218 | func AwaiterWithKey(ctx context.Context, key any) (completed func(), err error) { 219 | ctxx, completed := context.WithCancel(context.WithoutCancel(ctx)) //nolint:govet 220 | if err := CleanupWithKey(ctx, key, func() error { 221 | <-ctxx.Done() 222 | return nil 223 | }); err != nil { 224 | return nil, err //nolint:govet 225 | } 226 | return completed, nil 227 | } 228 | 229 | // Awaitable returns a function that guarantees execution of the process until it is called. 230 | // Note that if the timeout of WaitWithTimeout has passed (or the context of WaitWithContext has canceled), it will not wait. 231 | func Awaitable(ctx context.Context) (completed func()) { 232 | return AwaitableWithKey(ctx, doneGroupKey) 233 | } 234 | 235 | // AwaitableWithKey returns a function that guarantees execution of the process until it is called. 236 | // Note that if the timeout of WaitWithTimeout has passed (or the context of WaitWithContext has canceled), it will not wait. 237 | func AwaitableWithKey(ctx context.Context, key any) (completed func()) { 238 | completed, err := AwaiterWithKey(ctx, key) 239 | if err != nil { 240 | panic(err) 241 | } 242 | return completed 243 | } 244 | 245 | // Go calls the function now asynchronously. 246 | // If an error occurs, it is stored in the doneGroup. 247 | // Note that if the timeout of WaitWithTimeout has passed (or the context of WaitWithContext has canceled), it will not wait. 248 | func Go(ctx context.Context, f func() error) { 249 | GoWithKey(ctx, doneGroupKey, f) 250 | } 251 | 252 | // GoWithKey calls the function now asynchronously. 253 | // If an error occurs, it is stored in the doneGroup. 254 | // Note that if the timeout of WaitWithTimeout has passed (or the context of WaitWithContext has canceled), it will not wait. 255 | func GoWithKey(ctx context.Context, key any, f func() error) { 256 | dg, ok := ctx.Value(key).(*doneGroup) 257 | if !ok { 258 | panic(ErrNotContainDoneGroup) 259 | } 260 | completed, err := AwaiterWithKey(ctx, key) 261 | if err != nil { 262 | panic(err) 263 | } 264 | go func() { 265 | if err := f(); err != nil { 266 | dg.mu.Lock() 267 | dg.errors = errors.Join(dg.errors, err) 268 | dg.mu.Unlock() 269 | } 270 | completed() 271 | }() 272 | } 273 | 274 | func withDoneGroup(ctx context.Context, cancelCause context.CancelCauseFunc, key any) context.Context { 275 | wg := &sync.WaitGroup{} 276 | dg, ok := ctx.Value(key).(*doneGroup) 277 | if !ok { 278 | // Root doneGroup 279 | dg = &doneGroup{ 280 | cancel: cancelCause, 281 | cleanupGroups: []*sync.WaitGroup{wg}, 282 | } 283 | return context.WithValue(ctx, key, dg) 284 | } 285 | // Add cleanupGroup to parent doneGroup 286 | dg.cleanupGroups = append(dg.cleanupGroups, wg) 287 | 288 | // Leaf doneGroup 289 | leafDg := &doneGroup{ 290 | cancel: cancelCause, 291 | cleanupGroups: []*sync.WaitGroup{wg}, 292 | } 293 | return context.WithValue(ctx, key, leafDg) 294 | } 295 | -------------------------------------------------------------------------------- /donegroup_test.go: -------------------------------------------------------------------------------- 1 | package donegroup 2 | 3 | import ( 4 | "context" 5 | "errors" 6 | "sync/atomic" 7 | "testing" 8 | "time" 9 | ) 10 | 11 | func TestDoneGroup(t *testing.T) { 12 | t.Parallel() 13 | ctx, cancel := WithCancel(context.Background()) 14 | 15 | cleanup := atomic.Bool{} 16 | 17 | if err := Cleanup(ctx, func() error { 18 | time.Sleep(10 * time.Millisecond) 19 | cleanup.Store(true) 20 | return nil 21 | }); err != nil { 22 | t.Error(err) 23 | } 24 | 25 | defer func() { 26 | cancel() 27 | 28 | if err := Wait(ctx); err != nil { 29 | t.Error(err) 30 | } 31 | 32 | if !cleanup.Load() { 33 | t.Error("cleanup function not called") 34 | } 35 | }() 36 | 37 | cleanup.Store(false) 38 | } 39 | 40 | func TestCleanup(t *testing.T) { 41 | t.Parallel() 42 | t.Run("Cleanup with WithCancel", func(t *testing.T) { 43 | ctx, cancel := WithCancel(context.Background()) 44 | defer cancel() 45 | err := Cleanup(ctx, func() error { 46 | return nil 47 | }) 48 | if err != nil { 49 | t.Error(err) 50 | } 51 | }) 52 | 53 | t.Run("Cleanup without WithCancel", func(t *testing.T) { 54 | ctx := context.Background() 55 | err := Cleanup(ctx, func() error { 56 | return nil 57 | }) 58 | if !errors.Is(err, ErrNotContainDoneGroup) { 59 | t.Errorf("expected ErrNotContainDoneGroup, got %v", err) 60 | } 61 | }) 62 | } 63 | 64 | func TestWait(t *testing.T) { 65 | t.Parallel() 66 | t.Run("Wait with WithCancel", func(t *testing.T) { 67 | ctx, cancel := WithCancel(context.Background()) 68 | cancel() 69 | err := Wait(ctx) 70 | if err != nil { 71 | t.Error(err) 72 | } 73 | }) 74 | 75 | t.Run("Wait without WithCancel", func(t *testing.T) { 76 | ctx := context.Background() 77 | err := Wait(ctx) 78 | if !errors.Is(err, ErrNotContainDoneGroup) { 79 | t.Errorf("expected ErrNotContainDoneGroup, got %v", err) 80 | } 81 | }) 82 | 83 | t.Run("Collect errors", func(t *testing.T) { 84 | var ( 85 | errTest = errors.New("test error") 86 | errTest2 = errors.New("test error 2") 87 | ) 88 | 89 | ctx, cancel := WithCancel(context.Background()) 90 | if err := Cleanup(ctx, func() error { 91 | return errTest 92 | }); err != nil { 93 | t.Error(err) 94 | } 95 | if err := Cleanup(ctx, func() error { 96 | return errTest2 97 | }); err != nil { 98 | t.Error(err) 99 | } 100 | cancel() 101 | err := Wait(ctx) 102 | if !errors.Is(err, errTest) { 103 | t.Errorf("expected %v, got %v", errTest, err) 104 | } 105 | if !errors.Is(err, errTest2) { 106 | t.Errorf("expected %v, got %v", errTest2, err) 107 | } 108 | }) 109 | } 110 | 111 | func TestNoWait(t *testing.T) { 112 | t.Parallel() 113 | ctx, cancel := WithCancel(context.Background()) 114 | 115 | cleanup := atomic.Bool{} 116 | 117 | if err := Cleanup(ctx, func() error { 118 | time.Sleep(10 * time.Millisecond) 119 | cleanup.Store(true) 120 | return nil 121 | }); err != nil { 122 | t.Error(err) 123 | } 124 | 125 | defer func() { 126 | cancel() 127 | if cleanup.Load() { 128 | t.Error("cleanup function called") 129 | } 130 | 131 | time.Sleep(20 * time.Millisecond) 132 | if !cleanup.Load() { 133 | t.Error("cleanup function not called") 134 | } 135 | }() 136 | 137 | cleanup.Store(false) 138 | } 139 | 140 | func TestNoCleanup(t *testing.T) { 141 | t.Parallel() 142 | ctx, cancel := WithCancel(context.Background()) 143 | 144 | defer func() { 145 | cancel() 146 | 147 | if err := Wait(ctx); err != nil { 148 | t.Error(err) 149 | } 150 | }() 151 | } 152 | 153 | func TestMultiCleanup(t *testing.T) { 154 | t.Parallel() 155 | ctx, cancel := WithCancel(context.Background()) 156 | 157 | cleanup := atomic.Int64{} 158 | 159 | for i := 0; i < 10; i++ { 160 | if err := Cleanup(ctx, func() error { 161 | time.Sleep(10 * time.Millisecond) 162 | cleanup.Add(1) 163 | return nil 164 | }); err != nil { 165 | t.Error(err) 166 | } 167 | } 168 | 169 | defer func() { 170 | cancel() 171 | 172 | if err := Wait(ctx); err != nil { 173 | t.Error(err) 174 | } 175 | 176 | if cleanup.Load() != 10 { 177 | t.Error("cleanup function not called") 178 | } 179 | }() 180 | } 181 | 182 | func TestNestedWithCancel(t *testing.T) { 183 | t.Parallel() 184 | firstCtx, firstCancel := WithCancel(context.Background()) 185 | secondCtx, secondCancel := WithCancel(firstCtx) 186 | thirdCtx, thirdCancel := context.WithCancel(secondCtx) // context.WithCancel 187 | 188 | firstCleanup := atomic.Int64{} 189 | secondCleanup := atomic.Int64{} 190 | thirdCleanup := atomic.Int64{} 191 | 192 | for i := 0; i < 10; i++ { 193 | if err := Cleanup(firstCtx, func() error { 194 | time.Sleep(10 * time.Millisecond) 195 | firstCleanup.Add(1) 196 | return nil 197 | }); err != nil { 198 | t.Error(err) 199 | } 200 | } 201 | 202 | for i := 0; i < 5; i++ { 203 | if err := Cleanup(secondCtx, func() error { 204 | time.Sleep(10 * time.Millisecond) 205 | secondCleanup.Add(1) 206 | return nil 207 | }); err != nil { 208 | t.Error(err) 209 | } 210 | } 211 | 212 | for i := 0; i < 3; i++ { 213 | if err := Cleanup(thirdCtx, func() error { 214 | time.Sleep(10 * time.Millisecond) 215 | thirdCleanup.Add(1) 216 | return nil 217 | }); err != nil { 218 | t.Error(err) 219 | } 220 | } 221 | 222 | { 223 | dg, ok := firstCtx.Value(doneGroupKey).(*doneGroup) 224 | if !ok { 225 | t.Fatal("firstCtx.Value(doneGroupKey) is not *doneGroup") 226 | } 227 | got := len(dg.cleanupGroups) 228 | if want := 2; got != want { 229 | t.Errorf("firstCtx has %d cleanup groups, want %d", got, want) 230 | } 231 | } 232 | 233 | { 234 | dg, ok := secondCtx.Value(doneGroupKey).(*doneGroup) 235 | if !ok { 236 | t.Fatal("firstCtx.Value(doneGroupKey) is not *doneGroup") 237 | } 238 | got := len(dg.cleanupGroups) 239 | if want := 1; got != want { 240 | t.Errorf("firstCtx has %d cleanup groups, want %d", got, want) 241 | } 242 | } 243 | 244 | defer func() { 245 | thirdCancel() 246 | <-thirdCtx.Done() 247 | 248 | if firstCleanup.Load() != 0 { 249 | t.Error("cleanup function for first called") 250 | } 251 | if secondCleanup.Load() != 0 { 252 | t.Error("cleanup function for second called") 253 | } 254 | if thirdCleanup.Load() != 0 { 255 | t.Error("cleanup function for third called") 256 | } 257 | 258 | secondCancel() 259 | <-secondCtx.Done() 260 | 261 | if err := Wait(secondCtx); err != nil { 262 | t.Error(err) 263 | } 264 | 265 | if thirdCleanup.Load() != 3 { 266 | t.Error("cleanup function for third not called") 267 | } 268 | if secondCleanup.Load() != 5 { 269 | t.Error("cleanup function for second not called") 270 | } 271 | if firstCleanup.Load() != 0 { 272 | t.Error("cleanup function for first called") 273 | } 274 | 275 | firstCancel() 276 | <-firstCtx.Done() 277 | 278 | if err := Wait(firstCtx); err != nil { 279 | t.Error(err) 280 | } 281 | 282 | if thirdCleanup.Load() != 3 { 283 | t.Error("cleanup function for third not called") 284 | } 285 | if secondCleanup.Load() != 5 { 286 | t.Error("cleanup function for second not called") 287 | } 288 | if firstCleanup.Load() != 10 { 289 | t.Error("cleanup function for first not called") 290 | } 291 | }() 292 | } 293 | 294 | func TestRootWaitAll(t *testing.T) { 295 | t.Parallel() 296 | rootCtx, rootCancel := WithCancel(context.Background()) 297 | leafCtx, _ := WithCancel(rootCtx) 298 | 299 | rootCleanup := atomic.Int64{} 300 | leafCleanup := atomic.Int64{} 301 | 302 | for i := 0; i < 10; i++ { 303 | if err := Cleanup(rootCtx, func() error { 304 | time.Sleep(10 * time.Millisecond) 305 | rootCleanup.Add(1) 306 | return nil 307 | }); err != nil { 308 | t.Error(err) 309 | } 310 | } 311 | 312 | for i := 0; i < 5; i++ { 313 | if err := Cleanup(leafCtx, func() error { 314 | time.Sleep(10 * time.Millisecond) 315 | leafCleanup.Add(1) 316 | return nil 317 | }); err != nil { 318 | t.Error(err) 319 | } 320 | } 321 | 322 | defer func() { 323 | if rootCleanup.Load() != 0 { 324 | t.Error("cleanup function for root called") 325 | } 326 | 327 | rootCancel() 328 | 329 | if err := Wait(rootCtx); err != nil { 330 | t.Error(err) 331 | } 332 | 333 | if leafCleanup.Load() != 5 { 334 | t.Error("cleanup function for leaf not called") 335 | } 336 | 337 | if rootCleanup.Load() != 10 { 338 | t.Error("cleanup function for root not called") 339 | } 340 | }() 341 | } 342 | 343 | func TestWaitWithTimeout(t *testing.T) { 344 | t.Parallel() 345 | ctx, cancel := WithCancel(context.Background()) 346 | 347 | if err := Cleanup(ctx, func() error { 348 | for i := 0; i < 10; i++ { 349 | time.Sleep(2 * time.Millisecond) 350 | } 351 | return nil 352 | }); err != nil { 353 | t.Error(err) 354 | } 355 | 356 | timeout := 5 * time.Millisecond 357 | 358 | defer func() { 359 | cancel() 360 | time.Sleep(10 * time.Millisecond) 361 | if err := WaitWithTimeout(ctx, timeout); !errors.Is(err, context.DeadlineExceeded) { 362 | t.Error("expected timeout error") 363 | } 364 | }() 365 | } 366 | 367 | func TestWaitWithContext(t *testing.T) { 368 | t.Parallel() 369 | ctx, cancel := WithCancel(context.Background()) 370 | 371 | if err := Cleanup(ctx, func() error { 372 | for i := 0; i < 10; i++ { 373 | time.Sleep(2 * time.Millisecond) 374 | } 375 | return nil 376 | }); err != nil { 377 | t.Error(err) 378 | } 379 | 380 | timeout := 5 * time.Millisecond 381 | 382 | defer func() { 383 | cancel() 384 | ctxx, cancelx := context.WithTimeout(context.Background(), timeout) 385 | defer cancelx() 386 | time.Sleep(10 * time.Millisecond) 387 | if err := WaitWithContext(ctx, ctxx); !errors.Is(err, context.DeadlineExceeded) { 388 | t.Error("expected timeout error") 389 | } 390 | }() 391 | } 392 | 393 | func TestAwaiter(t *testing.T) { 394 | t.Parallel() 395 | tests := []struct { 396 | name string 397 | timeout time.Duration 398 | finished bool 399 | }{ 400 | { 401 | name: "finished", 402 | timeout: 100 * time.Millisecond, 403 | finished: true, 404 | }, 405 | { 406 | name: "not finished", 407 | timeout: 5 * time.Millisecond, 408 | finished: false, 409 | }, 410 | } 411 | for _, tt := range tests { 412 | tt := tt 413 | t.Run(tt.name, func(t *testing.T) { 414 | t.Parallel() 415 | ctx, cancel := WithCancel(context.Background()) 416 | 417 | var finished int32 418 | 419 | go func() { 420 | completed, err := Awaiter(ctx) 421 | if err != nil { 422 | t.Error(err) 423 | } 424 | <-ctx.Done() 425 | time.Sleep(20 * time.Millisecond) 426 | atomic.AddInt32(&finished, 1) 427 | completed() 428 | }() 429 | 430 | defer func() { 431 | cancel() 432 | time.Sleep(10 * time.Millisecond) 433 | err := WaitWithTimeout(ctx, tt.timeout) 434 | if tt.finished != (atomic.LoadInt32(&finished) > 0) { 435 | t.Errorf("expected finished: %v, got: %v", tt.finished, finished) 436 | } 437 | if tt.finished { 438 | if err != nil { 439 | t.Error(err) 440 | } 441 | return 442 | } 443 | if !errors.Is(err, context.DeadlineExceeded) { 444 | t.Errorf("expected timeout error: %v", err) 445 | } 446 | }() 447 | }) 448 | } 449 | } 450 | 451 | func TestAwaitable(t *testing.T) { 452 | t.Parallel() 453 | tests := []struct { 454 | name string 455 | timeout time.Duration 456 | finished bool 457 | }{ 458 | { 459 | name: "finished", 460 | timeout: 100 * time.Millisecond, 461 | finished: true, 462 | }, 463 | { 464 | name: "not finished", 465 | timeout: 5 * time.Millisecond, 466 | finished: false, 467 | }, 468 | } 469 | for _, tt := range tests { 470 | tt := tt 471 | t.Run(tt.name, func(t *testing.T) { 472 | t.Parallel() 473 | ctx, cancel := WithCancel(context.Background()) 474 | 475 | var finished int32 476 | 477 | go func() { 478 | defer Awaitable(ctx)() 479 | <-ctx.Done() 480 | time.Sleep(20 * time.Millisecond) 481 | atomic.AddInt32(&finished, 1) 482 | }() 483 | 484 | defer func() { 485 | cancel() 486 | time.Sleep(10 * time.Millisecond) 487 | err := WaitWithTimeout(ctx, tt.timeout) 488 | if tt.finished != (atomic.LoadInt32(&finished) > 0) { 489 | t.Errorf("expected finished: %v, got: %v", tt.finished, finished) 490 | } 491 | if tt.finished { 492 | if err != nil { 493 | t.Error(err) 494 | } 495 | return 496 | } 497 | if !errors.Is(err, context.DeadlineExceeded) { 498 | t.Errorf("expected timeout error: %v", err) 499 | } 500 | }() 501 | }) 502 | } 503 | } 504 | 505 | func TestCancel(t *testing.T) { 506 | t.Parallel() 507 | t.Run("Cancel with WithCancel", func(t *testing.T) { 508 | ctx, _ := WithCancel(context.Background()) 509 | err := Cancel(ctx) 510 | if err != nil { 511 | t.Error(err) 512 | } 513 | if !errors.Is(ctx.Err(), context.Canceled) { 514 | t.Error("expected context.Canceled") 515 | } 516 | }) 517 | 518 | t.Run("Cancel without WithCancel", func(t *testing.T) { 519 | ctx := context.Background() 520 | err := Cancel(ctx) 521 | if err == nil { 522 | t.Error("expected error, got nil") 523 | } 524 | }) 525 | } 526 | 527 | func TestGo(t *testing.T) { 528 | t.Parallel() 529 | tests := []struct { 530 | name string 531 | timeout time.Duration 532 | finished bool 533 | }{ 534 | { 535 | name: "finished", 536 | timeout: 200 * time.Millisecond, 537 | finished: true, 538 | }, 539 | { 540 | name: "not finished", 541 | timeout: 5 * time.Millisecond, 542 | finished: false, 543 | }, 544 | } 545 | for _, tt := range tests { 546 | tt := tt 547 | t.Run(tt.name, func(t *testing.T) { 548 | t.Parallel() 549 | ctx, cancel := WithCancel(context.Background()) 550 | 551 | var finished int32 552 | 553 | Go(ctx, func() error { 554 | <-ctx.Done() 555 | time.Sleep(100 * time.Millisecond) 556 | atomic.AddInt32(&finished, 1) 557 | return nil 558 | }) 559 | 560 | defer func() { 561 | cancel() 562 | time.Sleep(10 * time.Millisecond) 563 | err := WaitWithTimeout(ctx, tt.timeout) 564 | if tt.finished != (atomic.LoadInt32(&finished) > 0) { 565 | t.Errorf("expected finished: %v, got: %v", tt.finished, finished) 566 | } 567 | if tt.finished { 568 | if err != nil { 569 | t.Error(err) 570 | } 571 | return 572 | } 573 | if !errors.Is(err, context.DeadlineExceeded) { 574 | t.Errorf("expected timeout error: %v", err) 575 | } 576 | }() 577 | }) 578 | } 579 | } 580 | 581 | func TestGoWithError(t *testing.T) { 582 | t.Parallel() 583 | ctx, cancel := WithCancel(context.Background()) 584 | 585 | var errTest = errors.New("test error") 586 | 587 | Go(ctx, func() error { 588 | time.Sleep(10 * time.Millisecond) 589 | return errTest 590 | }) 591 | 592 | defer func() { 593 | cancel() 594 | 595 | err := Wait(ctx) 596 | if !errors.Is(err, errTest) { 597 | t.Errorf("got %v, want %v", err, errTest) 598 | } 599 | }() 600 | } 601 | 602 | func TestWithCancelCause(t *testing.T) { 603 | t.Parallel() 604 | ctx, cancel := WithCancelCause(context.Background()) 605 | 606 | cleanup := false 607 | 608 | if err := Cleanup(ctx, func() error { 609 | time.Sleep(10 * time.Millisecond) 610 | cleanup = true 611 | return nil 612 | }); err != nil { 613 | t.Error(err) 614 | } 615 | 616 | var errTest = errors.New("test error") 617 | 618 | defer func() { 619 | cancel(errTest) 620 | 621 | if err := Wait(ctx); err != nil { 622 | t.Error(err) 623 | } 624 | 625 | if !cleanup { 626 | t.Error("cleanup function not called") 627 | } 628 | 629 | if !errors.Is(context.Cause(ctx), errTest) { 630 | t.Errorf("got %v, want %v", context.Cause(ctx), errTest) 631 | } 632 | }() 633 | 634 | cleanup = false 635 | } 636 | 637 | func TestWithDeadline(t *testing.T) { 638 | t.Parallel() 639 | ctx, _ := WithDeadline(context.Background(), time.Now().Add(5*time.Millisecond)) 640 | 641 | cleanup := atomic.Bool{} 642 | 643 | if err := Cleanup(ctx, func() error { 644 | time.Sleep(10 * time.Millisecond) 645 | cleanup.Store(true) 646 | return nil 647 | }); err != nil { 648 | t.Error(err) 649 | } 650 | 651 | cleanup.Store(false) 652 | 653 | if err := Wait(ctx); err != nil { 654 | t.Error(err) 655 | } 656 | 657 | if !cleanup.Load() { 658 | t.Error("cleanup function not called") 659 | } 660 | 661 | if !errors.Is(context.Cause(ctx), context.DeadlineExceeded) { 662 | t.Errorf("got %v, want %v", context.Cause(ctx), context.DeadlineExceeded) 663 | } 664 | } 665 | 666 | func TestWithTimeout(t *testing.T) { 667 | t.Parallel() 668 | ctx, _ := WithTimeout(context.Background(), 5*time.Millisecond) 669 | 670 | cleanup := atomic.Bool{} 671 | 672 | if err := Cleanup(ctx, func() error { 673 | time.Sleep(10 * time.Millisecond) 674 | cleanup.Store(true) 675 | return nil 676 | }); err != nil { 677 | t.Error(err) 678 | } 679 | 680 | cleanup.Store(false) 681 | 682 | if err := Wait(ctx); err != nil { 683 | t.Error(err) 684 | } 685 | 686 | if !cleanup.Load() { 687 | t.Error("cleanup function not called") 688 | } 689 | 690 | if !errors.Is(context.Cause(ctx), context.DeadlineExceeded) { 691 | t.Errorf("got %v, want %v", context.Cause(ctx), context.DeadlineExceeded) 692 | } 693 | } 694 | 695 | func TestWithTimeoutCause(t *testing.T) { 696 | t.Parallel() 697 | var errTest = errors.New("test error") 698 | ctx, _ := WithTimeoutCause(context.Background(), 5*time.Millisecond, errTest) 699 | 700 | cleanup := atomic.Bool{} 701 | 702 | if err := Cleanup(ctx, func() error { 703 | time.Sleep(10 * time.Millisecond) 704 | cleanup.Store(true) 705 | return nil 706 | }); err != nil { 707 | t.Error(err) 708 | } 709 | 710 | cleanup.Store(false) 711 | 712 | if err := Wait(ctx); err != nil { 713 | t.Error(err) 714 | } 715 | 716 | if !cleanup.Load() { 717 | t.Error("cleanup function not called") 718 | } 719 | 720 | if !errors.Is(context.Cause(ctx), errTest) { 721 | t.Errorf("got %v, want %v", context.Cause(ctx), errTest) 722 | } 723 | } 724 | 725 | func TestCancelWithCause(t *testing.T) { 726 | t.Parallel() 727 | var errTest = errors.New("test error") 728 | var errTest2 = errors.New("test error2") 729 | 730 | t.Run("Cancel with cause", func(t *testing.T) { 731 | ctx, _ := WithCancel(context.Background()) 732 | 733 | if err := CancelWithCause(ctx, errTest); err != nil { 734 | t.Error(err) 735 | } 736 | 737 | if !errors.Is(context.Cause(ctx), errTest) { 738 | t.Errorf("got %v, want %v", context.Cause(ctx), errTest) 739 | } 740 | }) 741 | 742 | t.Run("Cancel with cause2", func(t *testing.T) { 743 | ctx, _ := WithCancel(context.Background()) 744 | 745 | if err := CancelWithCause(ctx, errTest); err != nil { 746 | t.Error(err) 747 | } 748 | 749 | if err := CancelWithCause(ctx, errTest2); err != nil { 750 | t.Error(err) 751 | } 752 | 753 | if !errors.Is(context.Cause(ctx), errTest) { 754 | t.Errorf("got %v, want %v", context.Cause(ctx), errTest) 755 | } 756 | if errors.Is(context.Cause(ctx), errTest2) { 757 | t.Error("got errTest2, want errTest") 758 | } 759 | }) 760 | } 761 | 762 | func TestWithoutCancel(t *testing.T) { 763 | t.Parallel() 764 | ctx, cancel := WithCancel(context.Background()) 765 | cleanup := atomic.Bool{} 766 | 767 | if err := Cleanup(ctx, func() error { 768 | cleanup.Store(true) 769 | return nil 770 | }); err != nil { 771 | t.Error(err) 772 | } 773 | 774 | ctx = WithoutCancel(ctx) 775 | 776 | if err := Wait(ctx); err == nil || !errors.Is(err, ErrNotContainDoneGroup) { 777 | t.Errorf("got %v, want %v", err, ErrNotContainDoneGroup) 778 | } 779 | 780 | cancel() 781 | 782 | time.Sleep(5 * time.Millisecond) 783 | 784 | if !cleanup.Load() { 785 | t.Error("cleanup function not called") 786 | } 787 | } 788 | -------------------------------------------------------------------------------- /example_test.go: -------------------------------------------------------------------------------- 1 | package donegroup_test 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "log" 7 | "time" 8 | 9 | "github.com/k1LoW/donegroup" 10 | ) 11 | 12 | func Example() { 13 | ctx, cancel := donegroup.WithCancel(context.Background()) 14 | 15 | // Cleanup process of some kind 16 | if err := donegroup.Cleanup(ctx, func() error { 17 | time.Sleep(10 * time.Millisecond) 18 | fmt.Println("cleanup with sleep") 19 | return nil 20 | }); err != nil { 21 | log.Fatal(err) 22 | } 23 | 24 | // Cleanup process of some kind 25 | if err := donegroup.Cleanup(ctx, func() error { 26 | fmt.Println("cleanup") 27 | return nil 28 | }); err != nil { 29 | log.Fatal(err) 30 | } 31 | 32 | defer func() { 33 | cancel() 34 | 35 | if err := donegroup.Wait(ctx); err != nil { 36 | log.Fatal(err) 37 | } 38 | }() 39 | 40 | // Main process of some kind 41 | fmt.Println("main start") 42 | 43 | fmt.Println("main finish") 44 | 45 | // Output: 46 | // main start 47 | // main finish 48 | // cleanup 49 | // cleanup with sleep 50 | } 51 | 52 | func ExampleAwaiter() { 53 | ctx, cancel := donegroup.WithCancel(context.Background()) 54 | 55 | go func() { 56 | completed, err := donegroup.Awaiter(ctx) 57 | if err != nil { 58 | log.Fatal(err) 59 | return 60 | } 61 | <-ctx.Done() 62 | time.Sleep(100 * time.Millisecond) 63 | fmt.Println("do something") 64 | completed() 65 | }() 66 | 67 | // Main process of some kind 68 | fmt.Println("main") 69 | time.Sleep(10 * time.Millisecond) 70 | 71 | cancel() 72 | if err := donegroup.Wait(ctx); err != nil { 73 | log.Fatal(err) 74 | } 75 | 76 | fmt.Println("finish") 77 | 78 | // Output: 79 | // main 80 | // do something 81 | // finish 82 | } 83 | 84 | func ExampleAwaitable() { 85 | ctx, cancel := donegroup.WithCancel(context.Background()) 86 | 87 | go func() { 88 | defer donegroup.Awaitable(ctx)() 89 | for { 90 | select { 91 | case <-ctx.Done(): 92 | time.Sleep(100 * time.Millisecond) 93 | fmt.Println("cleanup") 94 | return 95 | case <-time.After(10 * time.Millisecond): 96 | fmt.Println("do something") 97 | } 98 | } 99 | }() 100 | 101 | // Main process of some kind 102 | fmt.Println("main") 103 | time.Sleep(35 * time.Millisecond) 104 | 105 | cancel() 106 | if err := donegroup.Wait(ctx); err != nil { 107 | log.Fatal(err) 108 | } 109 | 110 | // Output: 111 | // main 112 | // do something 113 | // do something 114 | // do something 115 | // cleanup 116 | } 117 | 118 | func ExampleWaitWithTimeout() { 119 | ctx, cancel := donegroup.WithCancel(context.Background()) 120 | 121 | // Cleanup process of some kind 122 | if err := donegroup.Cleanup(ctx, func() error { 123 | fmt.Println("cleanup start") 124 | for i := 0; i < 10; i++ { 125 | time.Sleep(2 * time.Millisecond) 126 | } 127 | fmt.Println("cleanup finish") 128 | return nil 129 | }); err != nil { 130 | log.Fatal(err) 131 | } 132 | 133 | defer func() { 134 | cancel() 135 | timeout := 5 * time.Millisecond 136 | if err := donegroup.WaitWithTimeout(ctx, timeout); err != nil { 137 | fmt.Println(err) 138 | } 139 | }() 140 | 141 | // Main process of some kind 142 | fmt.Println("main start") 143 | 144 | fmt.Println("main finish") 145 | 146 | // Output: 147 | // main start 148 | // main finish 149 | // cleanup start 150 | // context deadline exceeded 151 | } 152 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/k1LoW/donegroup 2 | 3 | go 1.22.3 4 | -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/k1LoW/donegroup/21f315c3a0f6187aa6f2e09a1739832db27b4824/go.sum --------------------------------------------------------------------------------