Simple example of using go channels with context and having graceful termination

README.md

This gist provides a simple way to try out graceful termination

If any errors are there feel free to comment down below, I happy to fix them 👍

main.go

package main

import (
	"context"
	"errors"
	"log"
	"math/rand"
	"os/signal"
	"strconv"
	"sync"
	"syscall"
	"time"
)

type DB struct {
	wg   sync.WaitGroup
	lock sync.RWMutex
}

func (db *DB) Write() {
	db.lock.Lock()
	db.wg.Add(1)
	defer db.lock.Unlock()
	defer db.wg.Done()
	time.Sleep(time.Duration(rand.Intn(5)) * time.Second)
	log.Println("WRITE")
}

func (db *DB) Read() {
	db.lock.RLock()
	db.wg.Add(1)
	defer db.lock.RUnlock()
	defer db.wg.Done()
	time.Sleep(time.Duration(rand.Intn(2)) * time.Second)
	log.Println("READ")
}

type MainLoop struct {
	db *DB
}

func (m *MainLoop) Cleanup() {
	log.Println("[[STOP is triggered]]")
	m.db.Write()
	m.db.wg.Wait()
}

type Abcd struct {
	Name string
}

func main() {

	ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	ctx, cancel = context.WithTimeout(ctx, time.Duration(20*time.Second))
	resultsCh := make(chan *Abcd, 1)
	mainLoop := &MainLoop{db: &DB{}}
	i := 0
	for {
		select {
		case <-ctx.Done():
			if err := ctx.Err(); err != nil && !errors.Is(err, context.Canceled) {
				log.Printf("Reason of termination: %s\n", err)
			}
			log.Println("Context got cancelled")
			mainLoop.Cleanup()
			return
		case result := <-resultsCh:
			log.Printf("%#+v\n", result)
		default:
			if i > 20 {
				log.Println("normal termination")
				cancel()
				continue
			}
			log.Println("Waiting...", i)
			mainLoop.Work(resultsCh, i)
			i++
		}
	}
}

func (m *MainLoop) Work(data chan *Abcd, i int) {
	switch i {
	case 10, 20:
		data <- &Abcd{Name: "Dipankar" + strconv.Itoa(i)}
		m.db.Write()
	default:
		m.db.Read()
	}
}