Some example how to use context in Go

context-in-go.go

package main

import (
	"context"
	"log"
	"sync"
	"time"
)

type abTestVariant struct{}

func stage1(ctx context.Context) {
	ctx = context.WithValue(ctx, abTestVariant{}, "B")
	//ctx, cancel := context.WithCancel(ctx)
	//ctx, cancel := context.WithDeadline(ctx, time.Now().Add(time.Millisecond * 5000))
	ctx, cancel := context.WithTimeout(ctx, time.Millisecond * 5000)
	defer cancel()

	wg := sync.WaitGroup{}
	wg.Add(1)
	go func() {
		defer wg.Done()
		err := stage2(ctx)
		if err != nil {
			cancel()
		}
	}()
	wg.Add(1)
	go func() {
		defer wg.Done()
		err := stage3(ctx)
		if err != nil {
			cancel()
		}
	}()
	wg.Wait()
	if ctx.Err() != nil {
		log.Println("ERROR:", ctx.Err())
	} else {
		log.Println("DONE: STAGE 1")
	}
}

func stage2(ctx context.Context) error {
	variant, _ := ctx.Value(abTestVariant{}).(string)
	for i := 0;i < 10;i++ {
		select {
		case <- ctx.Done():
			log.Println("CANCELED: STAGE 2")
			return nil
		default:
			//if i == 2 {
			//	return fmt.Errorf("an error happened")
			//}
			time.Sleep(time.Millisecond * 100)
			log.Println("WORK: STAGE 2", variant)
		}
	}
	log.Println("DONE: STAGE 2")

	return nil
}

func stage3(ctx context.Context) error {
	for i := 0;i < 10;i++ {
		select {
		case <- ctx.Done():
			log.Println("CANCELED: STAGE 3")
			return nil
		default:
			time.Sleep(time.Millisecond * 100)
			log.Println("WORK: STAGE 3")
		}
	}
	log.Println("DONE: STAGE 3")

	return nil
}

func main() {
	ctx := context.Background()
	log.Println("START")
	stage1(ctx)
}