Golang sync/WaitGroup: new feature WaitWithTimeout

Hi! I would like to have a WaitWithTimeout function in the standard sync/WaitGroup structure.

Example:

package main

import (
    "context"
    "log"
    "os"
    "os/signal"
    "sync"
    "syscall"
    "time"
)

func main() {

    wg := &WaitGroupTimeout{}
    ctx, cancel := context.WithCancel(context.Background())

    var gracefulStopChan = make(chan os.Signal, 1)

    signal.Notify(gracefulStopChan, syscall.SIGTERM, syscall.SIGINT)

    go func() {
        sig := <-gracefulStopChan

        log.Printf("Caught sig: %+v\n", sig)
        log.Println("Application graceful shutdown begin...")

        // Shutdown

        cancel()
        if ok := wg.WaitWithTimeout(10 * time.Second); !ok {
            log.Println("Force shutdown by timeout")
        }
        log.Println("Application graceful shutdown finished")
        os.Exit(1)

    }()

    log.Println("APP start")

    wg.Add(1)
    go func() {
        defer wg.Done()
        for {
            select {
            case <-ctx.Done():
                //time.Sleep(4 * time.Second)
                for {
                    // failed shutdown process or closes something resource
                }
                log.Println("Break the process 1")
                return
            case <-time.After(1 * time.Second):
                log.Println("Process 1")
            }
        }
    }()

    wg.Add(1)
    go func() {
        defer wg.Done()
        for {
            select {
            case <-ctx.Done():
                time.Sleep(5 * time.Second)
                log.Println("Break the process 2")
                return
            case <-time.After(1 * time.Second):
                log.Println("Process 2")
            }
        }
    }()

    wg.Wait()
}

// I would not like to transfer this code from one project to another
type WaitGroupTimeout struct {
    sync.WaitGroup
}

// WaitWithTimeout returns the value "true" when  the [WaitGroup] counter is zero.
// And returns the value "false" when the wait is completed by timeout.
func (wg *WaitGroupTimeout) WaitWithTimeout(timeout time.Duration) bool {

    timeoutChan := time.After(timeout)
    waitChan := make(chan struct{})

    go func() {

        wg.Wait()
        close(waitChan)
    }()

    select {
    case <-timeoutChan:
        return false
    case <-waitChan:
        return true
    }
}

Comment From: gabyhelp

Related Issues and Documentation

• proposal: sync: add WaitContext to WaitGroup #40916
• In a concurrent scenario, will the mixed use of waitgroup and goroutine cause the program to be abnormal or hang up? #47665 (closed)
• sync: return true from the last WaitGroup.Done #7019 (closed)
• proposal: sync: add Go method to WaitGroup to launch a tracked goroutine #63796

_{(Emoji vote if this was helpful or unhelpful; more detailed feedback welcome in this discussion.)}

Comment From: seankhliao

Duplicate of #40916

Comment From: zigo101

It is easy to do this with an additional time.Timer.

Comment From: gozoro

I do this in one line wg.WaitWithTimeout(10 * time.Second). How to put your time.Timer into one function?

Comment From: zigo101

Certainly, you can't do it in one line. Programming with channels is fun, why to lose the fun? :D I mean it is fun to write the code in your first comment.