github.com/ydb-platform/ydb-go-sdk/v3@v3.89.2/internal/xcontext/cancels_quard.go (about)

     1  package xcontext
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  )
     7  
     8  type (
     9  	CancelsGuard struct {
    10  		mu      sync.Mutex
    11  		cancels map[*context.CancelFunc]struct{}
    12  	}
    13  )
    14  
    15  func NewCancelsGuard() *CancelsGuard {
    16  	return &CancelsGuard{
    17  		cancels: make(map[*context.CancelFunc]struct{}),
    18  	}
    19  }
    20  
    21  func (g *CancelsGuard) WithCancel(ctx context.Context) (context.Context, context.CancelFunc) {
    22  	g.mu.Lock()
    23  	defer g.mu.Unlock()
    24  	ctx, cancel := WithCancel(ctx)
    25  	g.cancels[&cancel] = struct{}{}
    26  
    27  	return ctx, func() {
    28  		cancel()
    29  		g.mu.Lock()
    30  		defer g.mu.Unlock()
    31  		delete(g.cancels, &cancel)
    32  	}
    33  }
    34  
    35  func (g *CancelsGuard) Cancel() {
    36  	g.mu.Lock()
    37  	defer g.mu.Unlock()
    38  	for cancel := range g.cancels {
    39  		(*cancel)()
    40  	}
    41  	g.cancels = make(map[*context.CancelFunc]struct{})
    42  }