diff --git a/syncs/syncs.go b/syncs/syncs.go new file mode 100644 index 0000000000000..17b73ca238305 --- /dev/null +++ b/syncs/syncs.go @@ -0,0 +1,66 @@ +// Copyright (c) 2020 Tailscale Inc & 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 syncs contains addition sync types. +package syncs + +import "sync/atomic" + +// ClosedChan returns a channel that's already closed. +func ClosedChan() <-chan struct{} { return closedChan } + +var closedChan = initClosedChan() + +func initClosedChan() <-chan struct{} { + ch := make(chan struct{}) + close(ch) + return ch +} + +// WaitGroupChan is like a sync.WaitGroup, but has a chan that closes +// on completion that you can wait on. (This, you can only use the +// value once) +// Also, its zero value is not usable. Use the constructor. +type WaitGroupChan struct { + n int64 // atomic + done chan struct{} // closed on transition to zero +} + +// NewWaitGroupChan returns a new single-use WaitGroupChan. +func NewWaitGroupChan() *WaitGroupChan { + return &WaitGroupChan{done: make(chan struct{})} +} + +// DoneChan returns a channel that's closed on completion. +func (c *WaitGroupChan) DoneChan() <-chan struct{} { return c.done } + +// Add adds delta, which may be negative, to the WaitGroupChan +// counter. If the counter becomes zero, all goroutines blocked on +// Wait or the Done chan are released. If the counter goes negative, +// Add panics. +// +// Note that calls with a positive delta that occur when the counter +// is zero must happen before a Wait. Calls with a negative delta, or +// calls with a positive delta that start when the counter is greater +// than zero, may happen at any time. Typically this means the calls +// to Add should execute before the statement creating the goroutine +// or other event to be waited for. +func (c *WaitGroupChan) Add(delta int) { + n := atomic.AddInt64(&c.n, int64(delta)) + if n == 0 { + close(c.done) + } +} + +// Decr decrements the WaitGroup counter by one. +// +// (It is like sync.WaitGroup's Done method, but we don't use Done in +// this type, because it's ambiguous between Context.Done and +// WaitGroup.Done. So we use DoneChan and Decr instead.) +func (wg *WaitGroupChan) Decr() { + wg.Add(-1) +} + +// Wait blocks until the WaitGroupChan counter is zero. +func (wg *WaitGroupChan) Wait() { <-wg.done } diff --git a/syncs/syncs_test.go b/syncs/syncs_test.go new file mode 100644 index 0000000000000..9de72e22f41fd --- /dev/null +++ b/syncs/syncs_test.go @@ -0,0 +1,50 @@ +// Copyright (c) 2020 Tailscale Inc & 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 syncs + +import "testing" + +func TestWaitGroupChan(t *testing.T) { + wg := NewWaitGroupChan() + + wantNotDone := func() { + t.Helper() + select { + case <-wg.DoneChan(): + t.Fatal("done too early") + default: + } + } + + wantDone := func() { + t.Helper() + select { + case <-wg.DoneChan(): + default: + t.Fatal("expected to be done") + } + } + + wg.Add(2) + wantNotDone() + + wg.Decr() + wantNotDone() + + wg.Decr() + wantDone() + wantDone() +} + +func TestClosedChan(t *testing.T) { + ch := ClosedChan() + for i := 0; i < 2; i++ { + select { + case <-ch: + default: + t.Fatal("not closed") + } + } +}