initial commit

Author: db7

LICENSE

@@ -0,0 +1,20 @@
+Copyright (c) db7
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+

barrier.go

@@ -0,0 +1,110 @@
+// Package barrier provides a data structure to synchronize a group of
+// goroutines, blocking them until all of the arrive on the barrier. Once the
+// last goroutine arrive an optional callback is executed in isolation.
+package barrier
+
+import (
+	"fmt"
+	"sync"
+)
+
+// ErrBarrierAborted is returned by Await() if Abort() was called.
+var ErrBarrierAborted = fmt.Errorf("Barrier aborted")
+
+// ErrBarrierMisused is returned by Await() if more than n concurrent Await()
+// calls are detected.
+var ErrBarrierMisused = fmt.Errorf("Barrier misused: more than n concurrent Await() calls")
+
+// Callback is called by the last goroutine entering the barrier.
+type Callback func() error
+
+// Barrier synchronizes a group of goroutines and optinally executes a callback
+// in isolation.
+type Barrier struct {
+	sync.Mutex
+	n     int64
+	count int64
+	done  chan bool
+	abort chan bool
+}
+
+// New returns a new Barrier which expects n goroutines to synchronize.
+func New(n int) *Barrier {
+	return &Barrier{
+		n:     int64(n),
+		count: int64(n),
+		done:  make(chan bool),
+		abort: make(chan bool),
+	}
+}
+
+// Abort marks the barrier as aborted and signal all waiting goroutines.
+// The barrier cannot be reset once aborted.
+func (b *Barrier) Abort() {
+	close(b.abort)
+}
+
+// Await synchronizes n goroutines and executes in isolation the callback of
+// the last goroutine calling Await. Await returns any error the callback
+// returns to one goroutine; if Abort() is called, ErrBarrierAborted is
+// returned. The number of goroutines call Await should always match the value
+// n passed in the barrier's initialization.
+func (b *Barrier) Await(cb Callback) error {
+	if b.aborted() {
+		return ErrBarrierAborted
+	}
+	// keep copy of current state
+	b.Lock()
+	b.count--
+	count := b.count
+	done := b.done
+	b.Unlock()
+
+	// more than n goroutines called Await
+	if count < 0 {
+		b.Abort()
+		return ErrBarrierMisused
+	}
+
+	// wait for others and callback execution
+	if count > 0 {
+		return b.wait(done)
+	}
+
+	// if count == 0 execute callback if last goroutine
+	var err error
+	if cb != nil {
+		err = cb()
+	}
+	b.reset()
+	return err
+}
+
+// aborted checks whether Barrier is aborted
+func (b *Barrier) aborted() bool {
+	select {
+	case <-b.abort:
+		return true
+	default:
+		return false
+	}
+}
+
+// wait waits for execution of callback or abort()
+func (b *Barrier) wait(done chan bool) error {
+	select {
+	case <-done:
+		return nil
+	case <-b.abort:
+		return ErrBarrierAborted
+	}
+}
+
+// reset resets the barrier for another round
+func (b *Barrier) reset() {
+	b.Lock()
+	close(b.done)
+	b.done = make(chan bool)
+	b.count = b.n
+	b.Unlock()
+}

barrier_test.go

@@ -0,0 +1,98 @@
+package barrier_test
+
+import (
+	"fmt"
+	"sync"
+	"sync/atomic"
+	"testing"
+
+	"github.com/db7/barrier"
+	"github.com/facebookgo/ensure"
+)
+
+func TestBarrier_manyrounds(t *testing.T) {
+	var count int64
+	rounds := 100
+	n := 10 // number of goroutines
+	b := barrier.New(n)
+
+	for i := 0; i < rounds; i++ {
+		var wg sync.WaitGroup
+		wg.Add(n)
+		for j := 0; j < n; j++ {
+			go func() {
+				b.Await(func() error {
+					atomic.AddInt64(&count, 1)
+					return nil
+				})
+				wg.Done()
+			}()
+		}
+		wg.Wait()
+	}
+	ensure.True(t, atomic.LoadInt64(&count) == int64(rounds))
+}
+
+func TestBarrier_abortBeforeLast(t *testing.T) {
+	n := 10 // number of goroutines
+	b := barrier.New(n)
+
+	// one round before abort
+	var wg sync.WaitGroup
+	wg.Add(n)
+	for j := 0; j < n; j++ {
+		go func() {
+			defer wg.Done()
+			err := b.Await(nil)
+			ensure.Nil(t, err)
+		}()
+	}
+	wg.Wait()
+
+	// one round when abort is called
+	wg.Add(n)
+	for j := 0; j < n-1; j++ {
+		go func() {
+			defer wg.Done()
+			err := b.Await(nil)
+			ensure.True(t, err == barrier.ErrBarrierAborted)
+		}()
+	}
+	// last goroutine
+	b.Abort()
+	go func() {
+		defer wg.Done()
+		err := b.Await(nil)
+		ensure.True(t, err == barrier.ErrBarrierAborted)
+	}()
+	wg.Wait()
+
+	// one last round where all goroutines should fail (aborted)
+	wg.Add(n)
+	for j := 0; j < n; j++ {
+		go func() {
+			defer wg.Done()
+			err := b.Await(nil)
+			ensure.True(t, err == barrier.ErrBarrierAborted)
+		}()
+	}
+	wg.Wait()
+}
+
+func ExampleBarrier_simple() {
+	n := 40 // number of goroutines
+	b := barrier.New(n)
+
+	var wg sync.WaitGroup
+	wg.Add(n)
+	for i := 0; i < n; i++ {
+		go func(k int) {
+			b.Await(func() error {
+				fmt.Println(k, "is the last goroutine")
+				return nil
+			})
+			wg.Done()
+		}(i)
+	}
+	wg.Wait()
+}