github.com/yaling888/clash@v1.53.0/common/singledo/singledo.go (about)

     1  package singledo
     2  
     3  import (
     4  	"sync"
     5  	"time"
     6  )
     7  
     8  type call[T any] struct {
     9  	wg  sync.WaitGroup
    10  	val T
    11  	err error
    12  }
    13  
    14  type Single[T any] struct {
    15  	mux    sync.Mutex
    16  	last   time.Time
    17  	wait   time.Duration
    18  	call   *call[T]
    19  	result *Result[T]
    20  }
    21  
    22  type Result[T any] struct {
    23  	Val T
    24  	Err error
    25  }
    26  
    27  // Do single.Do likes sync.singleFlight
    28  func (s *Single[T]) Do(fn func() (T, error)) (v T, err error, shared bool) {
    29  	s.mux.Lock()
    30  	now := time.Now()
    31  	if now.Before(s.last.Add(s.wait)) {
    32  		s.mux.Unlock()
    33  		return s.result.Val, s.result.Err, true
    34  	}
    35  
    36  	if callM := s.call; callM != nil {
    37  		s.mux.Unlock()
    38  		callM.wg.Wait()
    39  		return callM.val, callM.err, true
    40  	}
    41  
    42  	callM := &call[T]{}
    43  	callM.wg.Add(1)
    44  	s.call = callM
    45  	s.mux.Unlock()
    46  	callM.val, callM.err = fn()
    47  	callM.wg.Done()
    48  
    49  	s.mux.Lock()
    50  	s.call = nil
    51  	s.result = &Result[T]{callM.val, callM.err}
    52  	s.last = now
    53  	s.mux.Unlock()
    54  	return callM.val, callM.err, false
    55  }
    56  
    57  func (s *Single[T]) Reset() {
    58  	s.mux.Lock()
    59  	if s.result != nil {
    60  		var zero T
    61  		s.result.Val = zero
    62  		s.result.Err = nil
    63  		s.result = nil
    64  	}
    65  	s.last = time.Time{}
    66  	s.mux.Unlock()
    67  }
    68  
    69  func NewSingle[T any](wait time.Duration) *Single[T] {
    70  	return &Single[T]{wait: wait}
    71  }
    72  
    73  type Group[T any] struct {
    74  	mu sync.Mutex          // protects m
    75  	m  map[string]*call[T] // lazily initialized
    76  }
    77  
    78  func (g *Group[T]) Do(key string, fn func() (T, error)) (v T, err error, shared bool) {
    79  	g.mu.Lock()
    80  	if g.m == nil {
    81  		g.m = make(map[string]*call[T])
    82  	}
    83  	if c, ok := g.m[key]; ok {
    84  		g.mu.Unlock()
    85  		c.wg.Wait()
    86  
    87  		return c.val, c.err, true
    88  	}
    89  	c := new(call[T])
    90  	c.wg.Add(1)
    91  	g.m[key] = c
    92  	g.mu.Unlock()
    93  	c.val, c.err = fn()
    94  	c.wg.Done()
    95  
    96  	return c.val, c.err, false
    97  }
    98  
    99  func (g *Group[T]) Forget(key string) {
   100  	g.mu.Lock()
   101  	if g.m == nil {
   102  		g.mu.Unlock()
   103  		return
   104  	}
   105  	if c, ok := g.m[key]; ok {
   106  		var zero T
   107  		c.val = zero
   108  	}
   109  	delete(g.m, key)
   110  	g.mu.Unlock()
   111  }