CODEWITHSUNDEEP
go
septerr
septerr

Reputation: 6593

How to ensure context listener routine finishes before main exits?

Following is a simplified example of a scenario I am running into. Go Playground here.

The main method is periodically calling the process function. The process function acquires a lock that it must release before it returns or before the application is shut down due to an interrupt.

The interrupt is handled by cancelling the context passed to process. In my example, I am simply cancelling the context.

How can I ensure that when context is cancelled the unlocking logic is executed to completion?

ATM, in my tests, the logic is not getting executed to completion. The routine is kicked off, but does not seem to complete before the program exits.

It's as if the main method is exiting, killing the routine half way. Is that possible? How can I ensure that the routine always completes before the program exits? Thanks!

package main

import (
    "fmt"
    "context"
    "time"
    "sync"
)

func main() {
    ctx, cancel := context.WithCancel(context.Background())
    
    var wg sync.WaitGroup
    wg.Add(1)
    go func(){
        fmt.Println("Started periodic process")
        defer wg.Done()
        ticker := time.NewTicker(20 * time.Millisecond)
        for {
            select {
                case <-ticker.C:
                    process(ctx)
                case <-ctx.Done():
                    ticker.Stop()                   
    
            }
        }
        fmt.Println("Finished periodic process")
    }()
    
    time.Sleep(100 * time.Millisecond)  
    
    // cancel context
    cancel()
    
    wg.Wait()
}

func process(ctx context.Context) {
    ctx, cancel := context.WithCancel(ctx)
    defer cancel()
    
    go func() {
        <-ctx.Done()
        
        // unlock resources
        fmt.Println("Unlocking resources")

        time.Sleep(20 * time.Millisecond)

        fmt.Println("Unlocked resources")       
    }()
    
    // do some work
    fmt.Println("Started process")
    
    // acquire lock, in actual process
    // acquiring lock will fail unless
    // it was released be previous process
    // run
    fmt.Println("Acquired lock")
    
    time.Sleep(10 * time.Millisecond)
    
    fmt.Println("Finished process")

}

Upvotes: 0

Views: 142

Answers (1)

Brits
Brits

Reputation: 18370

As it stands your program will never exit. This is because the for loop in the goroutine started in main never exits (you stop the ticker when ctx.Done() but do not exit the loop).

A second issue is that in process the goroutine is cancelled when the function exits (via defer cancel()) but, due to the delay, the goroutine will continue to run for a period. The solution here depends upon whether you need the 'unlocking' to occur asynchronously (I have assumed this is not important).

The following (playground) resolves both issues and ensures that process does not return before the resources are free; if you need these to be freed in parallel then one option is to pass the WaitGroup to the function)

package main

import (
    "fmt"
    "context"
    "time"
    "sync"
)

func main() {
    ctx, cancel := context.WithCancel(context.Background())
    
    var wg sync.WaitGroup
    wg.Add(1)
    go func(){
            fmt.Println("Started periodic process")
        defer wg.Done()
        ticker := time.NewTicker(20 * time.Millisecond)
        for {
            select {
                case <-ticker.C:
                    process(ctx)
                case <-ctx.Done():
                    ticker.Stop()                   
                    fmt.Println("Finished periodic process")
                    return
            }
        }       
    }() 
    time.Sleep(100 * time.Millisecond)  
    cancel()    
    wg.Wait()
}

func process(ctx context.Context) {
    ctx, cancel := context.WithCancel(ctx)
    done := make(chan struct{}) // could also use a waitgroup
    
    go func() {
        <-ctx.Done()
        fmt.Println("Unlocking resources")
        time.Sleep(10 * time.Millisecond)
        fmt.Println("Unlocked resources")   
        close(done) 
    }()
    

    // do some work
    fmt.Println("Started process")  
    time.Sleep(10 * time.Millisecond)   
    fmt.Println("Finished process")
    
    cancel()
    <-done // Wait for resources to be unlocked
}

Upvotes: 1

Related Questions