Source file src/internal/singleflight/singleflight.go

     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 singleflight provides a duplicate function call suppression
     6	// mechanism.
     7	package singleflight
     8	
     9	import "sync"
    10	
    11	// call is an in-flight or completed singleflight.Do call
    12	type call struct {
    13		wg sync.WaitGroup
    14	
    15		// These fields are written once before the WaitGroup is done
    16		// and are only read after the WaitGroup is done.
    17		val interface{}
    18		err error
    19	
    20		// These fields are read and written with the singleflight
    21		// mutex held before the WaitGroup is done, and are read but
    22		// not written after the WaitGroup is done.
    23		dups  int
    24		chans []chan<- Result
    25	}
    26	
    27	// Group represents a class of work and forms a namespace in
    28	// which units of work can be executed with duplicate suppression.
    29	type Group struct {
    30		mu sync.Mutex       // protects m
    31		m  map[string]*call // lazily initialized
    32	}
    33	
    34	// Result holds the results of Do, so they can be passed
    35	// on a channel.
    36	type Result struct {
    37		Val    interface{}
    38		Err    error
    39		Shared bool
    40	}
    41	
    42	// Do executes and returns the results of the given function, making
    43	// sure that only one execution is in-flight for a given key at a
    44	// time. If a duplicate comes in, the duplicate caller waits for the
    45	// original to complete and receives the same results.
    46	// The return value shared indicates whether v was given to multiple callers.
    47	func (g *Group) Do(key string, fn func() (interface{}, error)) (v interface{}, err error, shared bool) {
    48		g.mu.Lock()
    49		if g.m == nil {
    50			g.m = make(map[string]*call)
    51		}
    52		if c, ok := g.m[key]; ok {
    53			c.dups++
    54			g.mu.Unlock()
    55			c.wg.Wait()
    56			return c.val, c.err, true
    57		}
    58		c := new(call)
    59		c.wg.Add(1)
    60		g.m[key] = c
    61		g.mu.Unlock()
    62	
    63		g.doCall(c, key, fn)
    64		return c.val, c.err, c.dups > 0
    65	}
    66	
    67	// DoChan is like Do but returns a channel that will receive the
    68	// results when they are ready. The second result is true if the function
    69	// will eventually be called, false if it will not (because there is
    70	// a pending request with this key).
    71	func (g *Group) DoChan(key string, fn func() (interface{}, error)) (<-chan Result, bool) {
    72		ch := make(chan Result, 1)
    73		g.mu.Lock()
    74		if g.m == nil {
    75			g.m = make(map[string]*call)
    76		}
    77		if c, ok := g.m[key]; ok {
    78			c.dups++
    79			c.chans = append(c.chans, ch)
    80			g.mu.Unlock()
    81			return ch, false
    82		}
    83		c := &call{chans: []chan<- Result{ch}}
    84		c.wg.Add(1)
    85		g.m[key] = c
    86		g.mu.Unlock()
    87	
    88		go g.doCall(c, key, fn)
    89	
    90		return ch, true
    91	}
    92	
    93	// doCall handles the single call for a key.
    94	func (g *Group) doCall(c *call, key string, fn func() (interface{}, error)) {
    95		c.val, c.err = fn()
    96		c.wg.Done()
    97	
    98		g.mu.Lock()
    99		delete(g.m, key)
   100		for _, ch := range c.chans {
   101			ch <- Result{c.val, c.err, c.dups > 0}
   102		}
   103		g.mu.Unlock()
   104	}
   105	
   106	// ForgetUnshared tells the singleflight to forget about a key if it is not
   107	// shared with any other goroutines. Future calls to Do for a forgotten key
   108	// will call the function rather than waiting for an earlier call to complete.
   109	// Returns whether the key was forgotten or unknown--that is, whether no
   110	// other goroutines are waiting for the result.
   111	func (g *Group) ForgetUnshared(key string) bool {
   112		g.mu.Lock()
   113		defer g.mu.Unlock()
   114		c, ok := g.m[key]
   115		if !ok {
   116			return true
   117		}
   118		if c.dups == 0 {
   119			delete(g.m, key)
   120			return true
   121		}
   122		return false
   123	}
   124	

View as plain text