Skip to content

Commit

Permalink
feat: introduce singleflight package with tests
Browse files Browse the repository at this point in the history
- Downgrade the Go version requirement from 1.21 to 1.18 in `go.mod`
- Add a new `singleflight` package to prevent duplicate function calls
- Implement `Group` and `call` types with methods `Do`, `DoChan`, `doCall`, and `ForgetUnshared` in `singleflight.go`
- Add tests for the `singleflight` package, including error handling and duplicate suppression

Signed-off-by: Bo-Yi Wu <[email protected]>
  • Loading branch information
appleboy committed Feb 4, 2024
1 parent 0e63ce0 commit 1dfacb0
Show file tree
Hide file tree
Showing 3 changed files with 200 additions and 1 deletion.
2 changes: 1 addition & 1 deletion example55-cache-hotspot-invalid/go.mod
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
module example55

go 1.21
go 1.18

require golang.org/x/sync v0.6.0
116 changes: 116 additions & 0 deletions example55-cache-hotspot-invalid/singleflight/singleflight.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,116 @@
// Copyright 2013 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// Package singleflight provides a duplicate function call suppression
// mechanism.
//
// Copy from https://go-review.googlesource.com/c/go/+/425187
package singleflight

import "sync"

// call is an in-flight or completed singleflight.Do call
type call[T any] struct {
wg sync.WaitGroup
// These fields are written once before the WaitGroup is done
// and are only read after the WaitGroup is done.
val T
err error
// These fields are read and written with the singleflight
// mutex held before the WaitGroup is done, and are read but
// not written after the WaitGroup is done.
dups int
chans []chan<- Result[T]
}

// Group represents a class of work and forms a namespace in
// which units of work can be executed with duplicate suppression.
type Group[K comparable, V any] struct {
mu sync.Mutex // protects m
m map[K]*call[V] // lazily initialized
}

// Result holds the results of Do, so they can be passed
// on a channel.
type Result[T any] struct {
Val T
Err error
Shared bool
}

// Do executes and returns the results of the given function, making
// sure that only one execution is in-flight for a given key at a
// time. If a duplicate comes in, the duplicate caller waits for the
// original to complete and receives the same results.
// The return value shared indicates whether v was given to multiple callers.
func (g *Group[K, V]) Do(key K, fn func() (V, error)) (v V, err error, shared bool) {
g.mu.Lock()
if g.m == nil {
g.m = make(map[K]*call[V])
}
if c, ok := g.m[key]; ok {
c.dups++
g.mu.Unlock()
c.wg.Wait()
return c.val, c.err, true
}
c := new(call[V])
c.wg.Add(1)
g.m[key] = c
g.mu.Unlock()
g.doCall(c, key, fn)
return c.val, c.err, c.dups > 0
}

// DoChan is like Do but returns a channel that will receive the
// results when they are ready.
func (g *Group[K, V]) DoChan(key K, fn func() (V, error)) <-chan Result[V] {
ch := make(chan Result[V], 1)
g.mu.Lock()
if g.m == nil {
g.m = make(map[K]*call[V])
}
if c, ok := g.m[key]; ok {
c.dups++
c.chans = append(c.chans, ch)
g.mu.Unlock()
return ch
}
c := &call[V]{chans: []chan<- Result[V]{ch}}
c.wg.Add(1)
g.m[key] = c
g.mu.Unlock()
go g.doCall(c, key, fn)
return ch
}

// doCall handles the single call for a key.
func (g *Group[K, V]) doCall(c *call[V], key K, fn func() (V, error)) {
c.val, c.err = fn()
c.wg.Done()
g.mu.Lock()
delete(g.m, key)
for _, ch := range c.chans {
ch <- Result[V]{c.val, c.err, c.dups > 0}
}
g.mu.Unlock()
}

// ForgetUnshared tells the singleflight to forget about a key if it is not
// shared with any other goroutines. Future calls to Do for a forgotten key
// will call the function rather than waiting for an earlier call to complete.
// Returns whether the key was forgotten or unknown--that is, whether no
// other goroutines are waiting for the result.
func (g *Group[K, V]) ForgetUnshared(key K) bool {
g.mu.Lock()
defer g.mu.Unlock()
c, ok := g.m[key]
if !ok {
return true
}
if c.dups == 0 {
delete(g.m, key)
return true
}
return false
}
83 changes: 83 additions & 0 deletions example55-cache-hotspot-invalid/singleflight/singleflight_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,83 @@
// Copyright 2013 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package singleflight

import (
"errors"
"fmt"
"sync"
"sync/atomic"
"testing"
"time"
)

func TestDo(t *testing.T) {
var g Group[string, any]
v, err, _ := g.Do("key", func() (any, error) {
return "bar", nil
})
if got, want := fmt.Sprintf("%v (%T)", v, v), "bar (string)"; got != want {
t.Errorf("Do = %v; want %v", got, want)
}
if err != nil {
t.Errorf("Do error = %v", err)
}
}

func TestDoErr(t *testing.T) {
var g Group[string, any]
someErr := errors.New("some error")
v, err, _ := g.Do("key", func() (any, error) {
return nil, someErr
})
if err != someErr {
t.Errorf("Do error = %v; want someErr %v", err, someErr)
}
if v != nil {
t.Errorf("unexpected non-nil value %#v", v)
}
}

func TestDoDupSuppress(t *testing.T) {
var g Group[string, any]
var wg1, wg2 sync.WaitGroup
c := make(chan string, 1)
var calls int32
fn := func() (any, error) {
if atomic.AddInt32(&calls, 1) == 1 {
// First invocation.
wg1.Done()
}
v := <-c
c <- v // pump; make available for any future calls
time.Sleep(10 * time.Millisecond) // let more goroutines enter Do
return v, nil
}
const n = 10
wg1.Add(1)
for i := 0; i < n; i++ {
wg1.Add(1)
wg2.Add(1)
go func() {
defer wg2.Done()
wg1.Done()
v, err, _ := g.Do("key", fn)
if err != nil {
t.Errorf("Do error: %v", err)
return
}
if s, _ := v.(string); s != "bar" {
t.Errorf("Do = %T %v; want %q", v, v, "bar")
}
}()
}
wg1.Wait()
// At least one goroutine is in fn now and all of them have at
// least reached the line before the Do.
c <- "bar"
wg2.Wait()
if got := atomic.LoadInt32(&calls); got <= 0 || got >= n {
t.Errorf("number of calls = %d; want over 0 and less than %d", got, n)
}
}

0 comments on commit 1dfacb0

Please sign in to comment.