Merge pull request #6 from arg0net:notifier-update

add StatefulNotifier.Update method

notifier.go

@@ -40,6 +40,20 @@ func (n *StatefulNotifier[T]) Load() (T, <-chan struct{}) {
 	return n.value, n.updated
 }
 
+// Update will atomically provide the current value to the update function
+// and store the result of the function.
+// Note that this will call the user's function with a lock held, so
+// if the function blocks, then other calls to the notifier will block.
+func (n *StatefulNotifier[T]) Update(fn func(T) T) {
+	n.mu.Lock()
+	defer n.mu.Unlock()
+
+	n.value = fn(n.value)
+	old := n.updated
+	n.updated = make(chan struct{})
+	close(old)
+}
+
 // Wait blocks until the given condition function returns true
 // or the context is canceled. It returns the value that satisfied the condition.
 //

notifier_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"math/rand"
 	"testing"
+	"sync"
 	"time"
 
 	"github.com/stretchr/testify/require"
@@ -29,6 +30,31 @@ func TestNotifier(t *testing.T) {
 	require.Equal(t, 4, v)
 }
 
+func TestNotifierUpdate(t *testing.T) {
+	sn := collections.NewStatefulNotifier(0)
+	start := make(chan struct{})
+
+	incr := func(in int) int {
+		return in+1
+	}
+	var wg sync.WaitGroup
+	for i := 0; i < 10; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			select {
+			case <-start:
+			}
+			sn.Update(incr)
+		}()
+	}
+	close(start)
+
+	wg.Wait()
+	v, _ := sn.Load()
+	require.Equal(t, 10, v)
+}
+
 func TestNotifierWait(t *testing.T) {
 	ctx := context.Background()