Merge pull request #137 from funny-falcon/selectable_future

futures/Selectable - future with selectable completion channel

Author: dustinhiatt-wf

futures/selectable.go

@@ -0,0 +1,105 @@
+/*
+Copyright 2016 Workiva, LLC
+Copyright 2016 Sokolov Yura aka funny_falcon
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package futures
+
+import (
+	"errors"
+	"sync"
+	"sync/atomic"
+)
+
+// ErrFutureCanceled signals that futures in canceled by a call to `f.Cancel()`
+var ErrFutureCanceled = errors.New("future canceled")
+
+// Selectable is a future with channel exposed for external `select`.
+// Many simultaneous listeners may wait for result either with `f.Value()`
+// or by selecting/fetching from `f.WaitChan()`, which is closed when future
+// fulfilled.
+type Selectable struct {
+	m      sync.Mutex
+	val    interface{}
+	err    error
+	wait   chan struct{}
+	filled uint32
+}
+
+// NewSelectable returns new selectable future.
+func NewSelectable() *Selectable {
+	return &Selectable{wait: make(chan struct{})}
+}
+
+func (f *Selectable) wchan() <-chan struct{} {
+	f.m.Lock()
+	ch := f.wait
+	f.m.Unlock()
+	return ch
+}
+
+// WaitChan returns channel, which is closed when future is fullfilled.
+func (f *Selectable) WaitChan() <-chan struct{} {
+	if atomic.LoadUint32(&f.filled) == 1 {
+		return closed
+	}
+	return f.wchan()
+}
+
+// GetResult waits for future to be fullfilled and returns value or error,
+// whatever is set first
+func (f *Selectable) GetResult() (interface{}, error) {
+	if atomic.LoadUint32(&f.filled) == 0 {
+		<-f.wchan()
+	}
+	return f.val, f.err
+}
+
+// Fill sets value for future, if it were not already fullfilled
+// Returns error, if it were already set to future.
+func (f *Selectable) Fill(v interface{}, e error) error {
+	f.m.Lock()
+	if f.filled == 0 {
+		f.val = v
+		f.err = e
+		atomic.StoreUint32(&f.filled, 1)
+		w := f.wait
+		f.wait = closed
+		close(w)
+	}
+	f.m.Unlock()
+	return f.err
+}
+
+// SetValue is alias for Fill(v, nil)
+func (f *Selectable) SetValue(v interface{}) error {
+	return f.Fill(v, nil)
+}
+
+// SetError is alias for Fill(nil, e)
+func (f *Selectable) SetError(e error) {
+	f.Fill(nil, e)
+}
+
+// Cancel is alias for SetError(ErrFutureCanceled)
+func (f *Selectable) Cancel() {
+	f.SetError(ErrFutureCanceled)
+}
+
+var closed = make(chan struct{})
+
+func init() {
+	close(closed)
+}

futures/selectable_test.go

@@ -0,0 +1,88 @@
+/*
+Copyright 2016 Workiva, LLC
+Copyright 2016 Sokolov Yura aka funny_falcon
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package futures
+
+import (
+	"fmt"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestSelectableGetResult(t *testing.T) {
+	f := NewSelectable()
+	var result interface{}
+	var err error
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		result, err = f.GetResult()
+		wg.Done()
+	}()
+
+	f.SetValue(`test`)
+	wg.Wait()
+
+	assert.Nil(t, err)
+	assert.Equal(t, `test`, result)
+
+	// ensure we don't get paused on the next iteration.
+	result, err = f.GetResult()
+
+	assert.Equal(t, `test`, result)
+	assert.Nil(t, err)
+}
+
+func TestSelectableSetError(t *testing.T) {
+	f := NewSelectable()
+	select {
+	case <-f.WaitChan():
+	case <-time.After(0):
+		f.SetError(fmt.Errorf("timeout"))
+	}
+
+	result, err := f.GetResult()
+
+	assert.Nil(t, result)
+	assert.NotNil(t, err)
+}
+
+func BenchmarkSelectable(b *testing.B) {
+	timeout := time.After(30 * time.Minute)
+	var wg sync.WaitGroup
+
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		wg.Add(1)
+		f := NewSelectable()
+		go func() {
+			select {
+			case <-f.WaitChan():
+			case <-timeout:
+				f.SetError(fmt.Errorf("timeout"))
+			}
+			wg.Done()
+		}()
+
+		f.SetValue(`test`)
+		wg.Wait()
+	}
+}