feat: simple future

Author: shils

async/async.go

@@ -0,0 +1,72 @@
+package async
+
+import (
+	"context"
+	"golang.org/x/sync/errgroup"
+)
+
+type Future[T any] interface {
+	Done() <-chan struct{}
+	Get() (T, error)
+}
+
+type future[T any] struct {
+	done  chan struct{}
+	value T
+	err   error
+}
+
+func (f *future[T]) Done() <-chan struct{} {
+	return f.done
+}
+
+func (f *future[T]) Get() (T, error) {
+	<-f.done
+	return f.value, f.err
+}
+
+func Compute[V any](fn func() (V, error)) Future[V] {
+	fut := &future[V]{done: make(chan struct{})}
+	go func() {
+		defer close(fut.done)
+		v, err := fn()
+		fut.value = v
+		fut.err = err
+	}()
+	return fut
+}
+
+func ComputeAll[K comparable, V any](ctx context.Context, fn func(context.Context, K) (V, error), keys ...K) Future[map[K]V] {
+	fut := &future[map[K]V]{done: make(chan struct{})}
+
+	numKeys := len(keys)
+	results := make([]V, numKeys)
+
+	g, gCtx := errgroup.WithContext(ctx)
+	g.SetLimit(numKeys)
+	go func() {
+		defer close(fut.done)
+		if err := g.Wait(); err != nil {
+			fut.err = err
+			return
+		}
+
+		fut.value = make(map[K]V)
+		for i, v := range results {
+			fut.value[keys[i]] = v
+		}
+	}()
+
+	for i, k := range keys {
+		index, key := i, k
+		g.Go(func() error {
+			v, err := fn(gCtx, key)
+			if err != nil {
+				return err
+			}
+			results[index] = v
+			return nil
+		})
+	}
+	return fut
+}

async/async_test.go

@@ -0,0 +1,91 @@
+package async
+
+import (
+	"context"
+	"errors"
+	"gotest.tools/v3/assert"
+	"testing"
+	"time"
+)
+
+func Test_ComputeError(t *testing.T) {
+	vCh := make(chan int)
+	eCh := make(chan error)
+	fut := Compute(func() (int, error) {
+		select {
+		case v := <-vCh:
+			return v, nil
+		case e := <-eCh:
+			return 0, e
+		}
+	})
+
+	e := errors.New("closed")
+	eCh <- e
+	res, err := fut.Get()
+	assert.Equal(t, res, 0)
+	assert.ErrorIs(t, err, e)
+}
+
+func Test_ComputeAll(t *testing.T) {
+	fut := ComputeAll(context.Background(), func(ctx context.Context, s string) (int, error) {
+		select {
+		case <-ctx.Done():
+			return 0, ctx.Err()
+		default:
+		}
+
+		time.Sleep(10 * time.Millisecond)
+		return len(s), nil
+	}, "a", "ab", "abc", "abcd")
+	expected := map[string]int{
+		"a":    1,
+		"ab":   2,
+		"abc":  3,
+		"abcd": 4,
+	}
+
+	result, err := fut.Get()
+	assert.NilError(t, err)
+	assert.DeepEqual(t, expected, result)
+}
+
+func Test_ComputeAllDone(t *testing.T) {
+	wait := make(chan struct{})
+	fut := ComputeAll(context.Background(), func(ctx context.Context, s string) (int, error) {
+		select {
+		case <-ctx.Done():
+			return 0, ctx.Err()
+		case <-wait:
+			return len(s), nil
+		}
+	}, "a", "ab", "abc", "abcd")
+
+	expected := map[string]int{
+		"a":    1,
+		"ab":   2,
+		"abc":  3,
+		"abcd": 4,
+	}
+
+	closed := false
+	for {
+		select {
+		case <-fut.Done():
+			result, err := fut.Get()
+			assert.NilError(t, err)
+			assert.DeepEqual(t, expected, result)
+			return
+		default:
+			if !closed {
+				close(wait)
+				closed = true
+			}
+		}
+	}
+}
+
+func zero[T any]() T {
+	var t T
+	return t
+}

go.mod

@@ -4,6 +4,7 @@ go 1.18
 
 require (
 	golang.org/x/exp v0.0.0-20230213192124-5e25df0256eb
+	golang.org/x/sync v0.1.0
 	gotest.tools/v3 v3.4.0
 )
 

go.sum

@@ -13,6 +13,8 @@ golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLL
 golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=
 golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.1.0 h1:wsuoTGHzEhffawBOhz5CYhcrV4IdKZbEyZjBMuTp12o=
+golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=