feat: implement futures to better handle async IPC & probes

cmd/probe.go

@@ -3,6 +3,7 @@ package cmd
 import (
 	"fmt"
 	"os"
+	"time"
 
 	"github.com/encodeous/nylon/core"
 	"github.com/encodeous/nylon/protocol"
@@ -17,8 +18,16 @@ var probeCmd = &cobra.Command{
 	Run: func(cmd *cobra.Command, args []string) {
 		itf, _ := cmd.Flags().GetString("interface")
 		jsonOut, _ := cmd.Flags().GetBool("json")
+		timeout, _ := cmd.Flags().GetDuration("timeout")
+		if timeout <= 0 {
+			fmt.Fprintln(os.Stderr, "Error: timeout must be positive")
+			os.Exit(1)
+		}
 		resp, err := core.SendIPCRequest(itf, &protocol.IpcRequest{
-			Request: &protocol.IpcRequest_Probe{Probe: &protocol.ProbeRequest{PeerId: args[0]}},
+			Request: &protocol.IpcRequest_Probe{Probe: &protocol.ProbeRequest{
+				PeerId:    args[0],
+				TimeoutMs: durationMillis(timeout),
+			}},
 		})
 		if err != nil {
 			fmt.Fprintln(os.Stderr, "Error:", err)
@@ -33,11 +42,14 @@ var probeCmd = &cobra.Command{
 			return
 		}
 		for _, r := range resp.GetProbe().Results {
-			status := "ok"
-			if !r.Success {
-				status = "error: " + r.Error
+			fmt.Printf("  %s  %s", r.Address, endpointProbeStatus(r))
+			if r.Resolved != nil && *r.Resolved != "" && *r.Resolved != r.Address {
+				fmt.Printf("  resolved=%s", *r.Resolved)
 			}
-			fmt.Printf("  %s  %s\n", r.Address, status)
+			if r.Status == protocol.EndpointProbeStatus_ENDPOINT_PROBE_REPLIED {
+				fmt.Printf("  latency=%s", formatDurationNs(r.LatencyNs))
+			}
+			fmt.Println()
 		}
 	},
 }
@@ -46,4 +58,31 @@ func init() {
 	rootCmd.AddCommand(probeCmd)
 	probeCmd.Flags().StringP("interface", "i", "nylon", "Interface name")
 	probeCmd.Flags().Bool("json", false, "Output as JSON")
+	probeCmd.Flags().Duration("timeout", 2*time.Second, "Probe response timeout")
+}
+
+func durationMillis(d time.Duration) uint32 {
+	ms := d.Milliseconds()
+	if ms <= 0 {
+		return 1
+	}
+	if ms > int64(^uint32(0)) {
+		return ^uint32(0)
+	}
+	return uint32(ms)
+}
+
+func endpointProbeStatus(r *protocol.EndpointProbeResult) string {
+	switch r.Status {
+	case protocol.EndpointProbeStatus_ENDPOINT_PROBE_REPLIED:
+		return "ok"
+	case protocol.EndpointProbeStatus_ENDPOINT_PROBE_TIMEOUT:
+		return "timeout"
+	case protocol.EndpointProbeStatus_ENDPOINT_PROBE_SEND_ERROR:
+		return "send-error"
+	case protocol.EndpointProbeStatus_ENDPOINT_PROBE_RESOLVE_ERROR:
+		return "resolve-error"
+	default:
+		return "unknown"
+	}
 }

core/future.go

@@ -0,0 +1,114 @@
+package core
+
+import (
+	"context"
+	"errors"
+	"sync"
+)
+
+var ErrFutureAlreadyAwaited = errors.New("future already awaited")
+
+type futureState[T any] struct {
+	done     chan struct{}
+	once     sync.Once
+	mu       sync.Mutex
+	value    T
+	err      error
+	awaiting bool
+	awaited  bool
+}
+
+// Future is a single-assignment async result.
+type Future[T any] struct {
+	state *futureState[T]
+}
+
+// NewFuture returns a Future and its completion function. Only the first
+// completion call is used; later calls are ignored.
+func NewFuture[T any]() (Future[T], func(T, error)) {
+	state := &futureState[T]{done: make(chan struct{})}
+	complete := func(value T, err error) {
+		state.once.Do(func() {
+			state.mu.Lock()
+			state.value = value
+			state.err = err
+			state.mu.Unlock()
+			close(state.done)
+		})
+	}
+	return Future[T]{state: state}, complete
+}
+
+// Await blocks until the future completes or ctx is canceled. Only one Await
+// call may wait for or consume a result at a time; concurrent calls and calls
+// after the result has been consumed return ErrFutureAlreadyAwaited.
+func (f Future[T]) Await(ctx context.Context) (T, error) {
+	if err := f.beginAwait(); err != nil {
+		var zero T
+		return zero, err
+	}
+	select {
+	case <-f.state.done:
+		return f.finishAwait()
+	default:
+	}
+	select {
+	case <-f.state.done:
+		return f.finishAwait()
+	case <-ctx.Done():
+		f.cancelAwait()
+		var zero T
+		return zero, ctx.Err()
+	}
+}
+
+func (f Future[T]) beginAwait() error {
+	f.state.mu.Lock()
+	defer f.state.mu.Unlock()
+	if f.state.awaited || f.state.awaiting {
+		return ErrFutureAlreadyAwaited
+	}
+	f.state.awaiting = true
+	return nil
+}
+
+func (f Future[T]) finishAwait() (T, error) {
+	f.state.mu.Lock()
+	defer f.state.mu.Unlock()
+	f.state.awaiting = false
+	f.state.awaited = true
+	return f.state.value, f.state.err
+}
+
+func (f Future[T]) cancelAwait() {
+	f.state.mu.Lock()
+	defer f.state.mu.Unlock()
+	f.state.awaiting = false
+}
+
+// Done returns a receive-only readiness channel for use in select statements.
+// Receive from Done, then call Await to consume the result.
+func (f Future[T]) Done() <-chan struct{} {
+	return f.state.done
+}
+
+// AwaitAll blocks until every input future completes or ctx is canceled.
+// Values are returned in the same order as the input slice. AwaitAll awaits and
+// consumes every completed input future in order.
+func AwaitAll[T any](ctx context.Context, futures []Future[T]) ([]T, error) {
+	values := make([]T, len(futures))
+	var firstErr error
+	for idx, item := range futures {
+		value, err := item.Await(ctx)
+		if err != nil {
+			if ctxErr := ctx.Err(); ctxErr != nil && errors.Is(err, ctxErr) {
+				return nil, ctxErr
+			}
+			if firstErr == nil {
+				firstErr = err
+			}
+		}
+		values[idx] = value
+	}
+	return values, firstErr
+}

core/future_test.go

@@ -0,0 +1,122 @@
+package core
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+)
+
+func TestFutureAwaitOnce(t *testing.T) {
+	future, complete := NewFuture[int]()
+	complete(7, nil)
+
+	value, err := future.Await(context.Background())
+	if err != nil {
+		t.Fatalf("first await failed: %v", err)
+	}
+	if value != 7 {
+		t.Fatalf("expected 7, got %d", value)
+	}
+
+	_, err = future.Await(context.Background())
+	if !errors.Is(err, ErrFutureAlreadyAwaited) {
+		t.Fatalf("expected ErrFutureAlreadyAwaited, got %v", err)
+	}
+}
+
+func TestFutureAwaitRejectsConcurrentWaiter(t *testing.T) {
+	future, complete := NewFuture[int]()
+	firstValue := make(chan int, 1)
+	firstErr := make(chan error, 1)
+
+	go func() {
+		value, err := future.Await(context.Background())
+		if err != nil {
+			firstErr <- err
+			return
+		}
+		firstValue <- value
+	}()
+
+	waitForFutureAwaiting(t, future)
+
+	_, err := future.Await(context.Background())
+	if !errors.Is(err, ErrFutureAlreadyAwaited) {
+		t.Fatalf("expected ErrFutureAlreadyAwaited, got %v", err)
+	}
+
+	complete(9, nil)
+	select {
+	case err := <-firstErr:
+		t.Fatalf("first await failed: %v", err)
+	case value := <-firstValue:
+		if value != 9 {
+			t.Fatalf("expected first await to get 9, got %d", value)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("timed out waiting for first await")
+	}
+}
+
+func TestFutureAwaitCancelDoesNotConsumeResult(t *testing.T) {
+	future, complete := NewFuture[int]()
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	_, err := future.Await(ctx)
+	if !errors.Is(err, context.Canceled) {
+		t.Fatalf("expected context canceled, got %v", err)
+	}
+
+	complete(11, nil)
+	value, err := future.Await(context.Background())
+	if err != nil {
+		t.Fatalf("second await failed: %v", err)
+	}
+	if value != 11 {
+		t.Fatalf("expected 11, got %d", value)
+	}
+}
+
+func TestFutureAllPreservesOrder(t *testing.T) {
+	a, completeA := NewFuture[int]()
+	b, completeB := NewFuture[int]()
+
+	completeB(2, nil)
+	completeA(1, nil)
+
+	values, err := AwaitAll(context.Background(), []Future[int]{a, b})
+	if err != nil {
+		t.Fatalf("await all failed: %v", err)
+	}
+	if len(values) != 2 || values[0] != 1 || values[1] != 2 {
+		t.Fatalf("unexpected values: %#v", values)
+	}
+}
+
+func TestFutureAllTimeout(t *testing.T) {
+	future, _ := NewFuture[int]()
+	ctx, cancel := context.WithTimeout(context.Background(), 0)
+	defer cancel()
+
+	_, err := AwaitAll(ctx, []Future[int]{future})
+	if !errors.Is(err, context.DeadlineExceeded) {
+		t.Fatalf("expected context deadline exceeded, got %v", err)
+	}
+}
+
+func waitForFutureAwaiting[T any](t *testing.T, future Future[T]) {
+	t.Helper()
+	deadline := time.Now().Add(time.Second)
+	for time.Now().Before(deadline) {
+		future.state.mu.Lock()
+		awaiting := future.state.awaiting
+		future.state.mu.Unlock()
+		if awaiting {
+			return
+		}
+		time.Sleep(time.Millisecond)
+	}
+	t.Fatal("timed out waiting for future to have a waiter")
+}