Simplify PublicationAwaiter (#718)

Make the `PublicationAwaiter` more concise by using a `sync.Cond` so that calls to `Await()` are notified to check and release themselves when an updated checkpoint is seen, rather than relying on managing many channels.

Author: AlCutter

await.go

@@ -17,13 +17,10 @@ package tessera
 import (
 	"context"
 	"errors"
-	"fmt"
 	"os"
 	"sync"
 	"time"
 
-	"container/list"
-
 	"github.com/transparency-dev/tessera/internal/parse"
 	"k8s.io/klog/v2"
 )
@@ -33,7 +30,7 @@ import (
 // to fetch checkpoints using the `readCheckpoint` function.
 func NewPublicationAwaiter(ctx context.Context, readCheckpoint func(ctx context.Context) ([]byte, error), pollPeriod time.Duration) *PublicationAwaiter {
 	a := &PublicationAwaiter{
-		waiters: list.New(),
+		c: sync.NewCond(&sync.Mutex{}),
 	}
 	go a.pollLoop(ctx, readCheckpoint, pollPeriod)
 	return a
@@ -54,24 +51,13 @@ func NewPublicationAwaiter(ctx context.Context, readCheckpoint func(ctx context.
 // When used this way, it requires very little code at the point of use to
 // block until the new leaf is integrated into the tree.
 type PublicationAwaiter struct {
-	// This mutex is used to protect all reads and writes to fields below.
-	mu sync.Mutex
-	// waiters is a linked list of type `waiter` and corresponds to all of the
-	// client threads that are currently blocked on this awaiter. The list is
-	// not sorted; new clients are simply added to the end. This can be changed
-	// if this is the wrong decision, but the reasoning behind this decision is
-	// that it is O(1) at the point of adding an entry to simply put it on the
-	// end, and O(N) each time we poll to check all clients. Keeping the list
-	// sorted by index would reduce the worst case in the poll loop, but increase
-	// the cost at the point of adding a new entry.
-	//
-	// Once a waiter is added to this list, it belongs to the awaiter and the
-	// awaiter takes sole responsibility for closing the channel in the waiter.
-	waiters *list.List
-	// size and checkpoint keep track of the latest seen size and checkpoint as
-	// an optimization where the last seen value was already large enough.
+	c *sync.Cond
+
+	// size, checkpoint, and err keep track of the latest size and checkpoint
+	// (or error) seen by the poller.
 	size       uint64
 	checkpoint []byte
+	err        error
 }
 
 // Await blocks until the IndexFuture is resolved, and this new index has been
@@ -83,15 +69,24 @@ type PublicationAwaiter struct {
 // or in the event that there is an error getting a valid checkpoint, an error
 // will be returned from this method.
 func (a *PublicationAwaiter) Await(ctx context.Context, future IndexFuture) (Index, []byte, error) {
-	ctx, span := tracer.Start(ctx, "tessera.Await")
+	_, span := tracer.Start(ctx, "tessera.Await")
 	defer span.End()
 
 	i, err := future()
 	if err != nil {
 		return i, nil, err
 	}
-	cp, err := a.await(ctx, i.Index)
-	return i, cp, err
+
+	a.c.L.Lock()
+	defer a.c.L.Unlock()
+	for (a.size <= i.Index && a.err == nil) && ctx.Err() == nil {
+		a.c.Wait()
+	}
+	// Ensure we propogate context done error, if any.
+	if err := ctx.Err(); err != nil {
+		a.err = err
+	}
+	return i, a.checkpoint, a.err
 }
 
 // pollLoop MUST be called in a goroutine when constructing an PublicationAwaiter
@@ -100,101 +95,36 @@ func (a *PublicationAwaiter) Await(ctx context.Context, future IndexFuture) (Ind
 // the latest checkpoint from the log, parses the tree size, and releases all clients
 // that were blocked on an index smaller than this tree size.
 func (a *PublicationAwaiter) pollLoop(ctx context.Context, readCheckpoint func(ctx context.Context) ([]byte, error), pollPeriod time.Duration) {
-	for {
+	var (
+		cp     []byte
+		cpErr  error
+		cpSize uint64
+	)
+	for done := false; !done; {
 		select {
 		case <-ctx.Done():
 			klog.Info("PublicationAwaiter exiting due to context completion")
-			return
+			cp, cpSize, cpErr = nil, 0, ctx.Err()
+			done = true
 		case <-time.After(pollPeriod):
-			// It's worth this small lock contention to make sure that no unnecessary
-			// work happens in personalities that aren't performing writes.
-			a.mu.Lock()
-			hasClients := a.waiters.Front() != nil
-			a.mu.Unlock()
-			if !hasClients {
+			cp, cpErr = readCheckpoint(ctx)
+			switch {
+			case errors.Is(cpErr, os.ErrNotExist):
 				continue
+			case cpErr != nil:
+				cpSize = 0
+			default:
+				_, cpSize, _, cpErr = parse.CheckpointUnsafe(cp)
 			}
 		}
+
+		a.c.L.Lock()
 		// Note that for now, this releases all clients in the event of a single failure.
 		// If this causes problems, this could be changed to attempt retries.
-		rawCp, err := readCheckpoint(ctx)
-		if err != nil && !errors.Is(err, os.ErrNotExist) {
-			a.releaseClientsErr(fmt.Errorf("readCheckpoint: %v", err))
-			continue
-		}
-		_, size, _, err := parse.CheckpointUnsafe(rawCp)
-		if err != nil {
-			a.releaseClientsErr(err)
-			continue
-		}
-		a.releaseClients(size, rawCp)
-	}
-}
-
-func (a *PublicationAwaiter) await(ctx context.Context, i uint64) ([]byte, error) {
-	a.mu.Lock()
-	if a.size > i {
-		cp := a.checkpoint
-		a.mu.Unlock()
-		return cp, nil
-	}
-	done := make(chan checkpointOrDeath, 1)
-	w := waiter{
-		index:  i,
-		result: done,
-	}
-	a.waiters.PushBack(w)
-	a.mu.Unlock()
-	select {
-	case <-ctx.Done():
-		return nil, ctx.Err()
-	case cod := <-done:
-		return cod.cp, cod.err
+		a.checkpoint = cp
+		a.size = cpSize
+		a.err = cpErr
+		a.c.Broadcast()
+		a.c.L.Unlock()
 	}
 }
-
-func (a *PublicationAwaiter) releaseClientsErr(err error) {
-	cod := checkpointOrDeath{
-		err: err,
-	}
-	a.mu.Lock()
-	defer a.mu.Unlock()
-	for e := a.waiters.Front(); e != nil; e = e.Next() {
-		w := e.Value.(waiter)
-		w.result <- cod
-		close(w.result)
-	}
-	// Clear out the list now we've released everyone
-	a.waiters.Init()
-}
-
-func (a *PublicationAwaiter) releaseClients(size uint64, cp []byte) {
-	cod := checkpointOrDeath{
-		cp: cp,
-	}
-	a.mu.Lock()
-	defer a.mu.Unlock()
-	a.size = size
-	a.checkpoint = cp
-	for e := a.waiters.Front(); e != nil; e = e.Next() {
-		w := e.Value.(waiter)
-		if w.index < size {
-			w.result <- cod
-			close(w.result)
-			// Need to do this removal after the loop has been fully iterated
-			// It still needs to happen inside the mutex, but defers happen in
-			// reverse order.
-			defer a.waiters.Remove(e)
-		}
-	}
-}
-
-type waiter struct {
-	index  uint64
-	result chan checkpointOrDeath
-}
-
-type checkpointOrDeath struct {
-	cp  []byte
-	err error
-}

await_test.go

@@ -12,21 +12,21 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package tessera_test
+package tessera
 
 import (
 	"bytes"
 	"context"
 	"crypto/sha256"
 	"fmt"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"errors"
 	"testing"
 
 	"github.com/transparency-dev/formats/log"
-	"github.com/transparency-dev/tessera"
 	"golang.org/x/mod/sumdb/note"
 )
 
@@ -136,11 +136,11 @@ func TestAwait(t *testing.T) {
 				<-time.After(tC.cpDelay)
 				return tC.cpBody, tC.cpErr
 			}
-			awaiter := tessera.NewPublicationAwaiter(ctx, readCheckpoint, 10*time.Millisecond)
+			awaiter := NewPublicationAwaiter(ctx, readCheckpoint, 10*time.Millisecond)
 
-			future := func() (tessera.Index, error) {
+			future := func() (Index, error) {
 				<-time.After(tC.fDelay)
-				return tessera.Index{Index: tC.fIndex}, tC.fErr
+				return Index{Index: tC.fIndex}, tC.fErr
 			}
 			i, cp, err := awaiter.Await(ctx, future)
 			if gotErr := err != nil; gotErr != tC.wantErr {
@@ -195,14 +195,14 @@ func TestAwait_multiClient(t *testing.T) {
 		}
 		return n, nil
 	}
-	awaiter := tessera.NewPublicationAwaiter(ctx, readCheckpoint, 10*time.Millisecond)
+	awaiter := NewPublicationAwaiter(ctx, readCheckpoint, 10*time.Millisecond)
 
 	wg := sync.WaitGroup{}
 	for i := range 300 {
 		index := uint64(i)
-		future := func() (tessera.Index, error) {
+		future := func() (Index, error) {
 			<-time.After(15 * time.Millisecond)
-			return tessera.Index{Index: index}, nil
+			return Index{Index: index}, nil
 		}
 		wg.Add(1)
 		go func() {
@@ -226,3 +226,29 @@ func TestAwait_multiClient(t *testing.T) {
 	}
 	wg.Wait()
 }
+
+func BenchmarkAwait(b *testing.B) {
+	cpFormat := "origin/\n%d\nhash\n\nsig"
+	cpSize := atomic.Uint64{}
+	readCP := func(_ context.Context) ([]byte, error) {
+		return fmt.Appendf(nil, cpFormat, cpSize.Load()), nil
+	}
+	a := NewPublicationAwaiter(b.Context(), readCP, time.Millisecond)
+	t := time.NewTicker(time.Millisecond)
+	go func() {
+		for {
+			select {
+			case <-b.Context().Done():
+				return
+			case <-t.C:
+				cpSize.Add(1)
+			}
+		}
+	}()
+	f := func() (Index, error) {
+		return Index{Index: cpSize.Load()}, nil
+	}
+	for b.Loop() {
+		_, _, _ = a.Await(b.Context(), f)
+	}
+}