github.com/phuslu/lru@v1.0.16-0.20240421170520-46288a2fd47c/singleflight.go (about)

     1  // Copyright 2013 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package lru
     6  
     7  import "sync"
     8  
     9  // singleflight_call is an in-flight or completed singleflight.Do singleflight_call
    10  type singleflight_call[T any] struct {
    11  	wg sync.WaitGroup
    12  
    13  	// These fields are written once before the WaitGroup is done
    14  	// and are only read after the WaitGroup is done.
    15  	val T
    16  	err error
    17  
    18  	// These fields are read and written with the singleflight
    19  	// mutex held before the WaitGroup is done, and are read but
    20  	// not written after the WaitGroup is done.
    21  	dups int
    22  }
    23  
    24  // Group represents a class of work and forms a namespace in
    25  // which units of work can be executed with duplicate suppression.
    26  type singleflight_Group[K comparable, V any] struct {
    27  	mu sync.Mutex                  // protects m
    28  	m  map[K]*singleflight_call[V] // lazily initialized
    29  }
    30  
    31  // Result holds the results of Do, so they can be passed
    32  // on a channel.
    33  type singleflight_Result[T any] struct {
    34  	Val    T
    35  	Err    error
    36  	Shared bool
    37  }
    38  
    39  // Do executes and returns the results of the given function, making
    40  // sure that only one execution is in-flight for a given key at a
    41  // time. If a duplicate comes in, the duplicate singleflight_caller waits for the
    42  // original to complete and receives the same results.
    43  // The return value shared indicates whether v was given to multiple singleflight_callers.
    44  func (g *singleflight_Group[K, V]) Do(key K, fn func() (V, error)) (v V, err error, shared bool) {
    45  	g.mu.Lock()
    46  	if g.m == nil {
    47  		g.m = make(map[K]*singleflight_call[V])
    48  	}
    49  	if c, ok := g.m[key]; ok {
    50  		c.dups++
    51  		g.mu.Unlock()
    52  		c.wg.Wait()
    53  		return c.val, c.err, true
    54  	}
    55  	c := new(singleflight_call[V])
    56  	c.wg.Add(1)
    57  	g.m[key] = c
    58  	g.mu.Unlock()
    59  
    60  	g.doCall(c, key, fn)
    61  	return c.val, c.err, c.dups > 0
    62  }
    63  
    64  // doCall handles the single singleflight_call for a key.
    65  func (g *singleflight_Group[K, V]) doCall(c *singleflight_call[V], key K, fn func() (V, error)) {
    66  	c.val, c.err = fn()
    67  	c.wg.Done()
    68  
    69  	g.mu.Lock()
    70  	delete(g.m, key)
    71  	g.mu.Unlock()
    72  }