feat: implement shared watch stream

pkg/state/impl/etcd/controller_runtime_test.go

@@ -36,10 +36,32 @@ func init() {
 func TestRuntimeConformance(t *testing.T) {
 	t.Parallel()
 
+	runtimeConformance(t)
+}
+
+// TestRuntimeConformanceSharedWatch runs the controller runtime conformance suite against a State
+// serving all of its watches from a single shared etcd watcher.
+//
+// The controller runtime funnels every one of its per-type watches into a single channel, so it is
+// the consumer which depends on the cross-kind ordering the shared watcher provides.
+func TestRuntimeConformanceSharedWatch(t *testing.T) {
+	t.Parallel()
+
+	runtimeConformance(t, etcd.WithSharedWatch())
+}
+
+func runtimeConformance(t *testing.T, opts ...etcd.StateOption) {
+	t.Helper()
+
 	testhelpers.WithEtcd(t, func(cli *clientv3.Client) {
 		suite := &conformance.RuntimeSuite{
 			SetupRuntime: func(rs *conformance.RuntimeSuite) {
-				etcdState := etcd.NewState(cli, store.ProtobufMarshaler{}, etcd.WithSalt([]byte("test123")), etcd.WithKeyPrefix(rs.T().Name()))
+				stateOpts := append([]etcd.StateOption{
+					etcd.WithSalt([]byte("test123")),
+					etcd.WithKeyPrefix(rs.T().Name()),
+				}, opts...)
+
+				etcdState := etcd.NewState(cli, store.ProtobufMarshaler{}, stateOpts...)
 				rs.State = state.WrapCore(etcdState)
 				rs.Runtime = tmust.Value(runtime.NewRuntime(rs.State, logging.DefaultLogger()))(rs.T())
 			},

pkg/state/impl/etcd/etcd.go

@@ -40,6 +40,7 @@ type State struct {
 	marshaler store.Marshaler
 	observer  ObserverFunc
 	limiter   LimiterFunc
+	mux       *watchMux
 	keyPrefix string
 	salt      []byte
 }
@@ -55,14 +56,20 @@ func NewState(cli Client, marshaler store.Marshaler, opts ...StateOption) *State
 		opt(&options)
 	}
 
-	return &State{
+	st := &State{
 		cli:       cli,
 		marshaler: marshaler,
 		observer:  options.observer,
 		limiter:   options.limiter,
 		keyPrefix: options.keyPrefix,
 		salt:      options.salt,
 	}
+
+	if options.sharedWatch {
+		st.mux = &watchMux{st: st}
+	}
+
+	return st
 }
 
 func (st *State) invokeHook(
@@ -430,6 +437,20 @@ func (st *State) Watch(ctx context.Context, resourcePointer resource.Pointer, ch
 
 	etcdKey := st.etcdKeyFromPointer(resourcePointer)
 
+	// watches resuming from a bookmark are positioned before the shared watcher, which cannot
+	// replay history, so they keep using a dedicated etcd watcher
+	if st.mux != nil && options.TailEvents == 0 && options.StartFromBookmark == nil {
+		return st.watchMuxed(ctx, &subscriber{
+			ctx:      ctx,
+			st:       st,
+			mux:      st.mux,
+			pointer:  resourcePointer,
+			key:      etcdKey,
+			exact:    true,
+			singleCh: ch,
+		}, "watch")
+	}
+
 	var (
 		revision     int64
 		initialEvent state.Event
@@ -577,6 +598,21 @@ func (st *State) watchKind(ctx context.Context, resourceKind resource.Kind, sing
 
 	etcdKey := st.etcdKeyPrefixFromKind(resourceKind)
 
+	// watches resuming from a bookmark are positioned before the shared watcher, which cannot
+	// replay history, so they keep using a dedicated etcd watcher
+	if st.mux != nil && options.TailEvents == 0 && options.StartFromBookmark == nil {
+		return st.watchMuxed(ctx, &subscriber{
+			ctx:      ctx,
+			st:       st,
+			mux:      st.mux,
+			kind:     resourceKind,
+			options:  options,
+			key:      etcdKey,
+			singleCh: singleCh,
+			aggCh:    aggCh,
+		}, opName)
+	}
+
 	var (
 		bootstrapList []resource.Resource
 		revision      int64

pkg/state/impl/etcd/key.go

@@ -8,6 +8,7 @@ import (
 	"crypto/sha256"
 	"encoding/hex"
 	"net/url"
+	"strings"
 
 	"github.com/cosi-project/runtime/pkg/resource"
 )
@@ -29,6 +30,20 @@ func (st *State) etcdKeyPrefixFromKind(kind resource.Kind) string {
 	return st.keyPrefix + "/" + nsEscaped + "/" + typeEscaped + "/"
 }
 
+// etcdKeyPrefixFromKey returns the kind prefix of the given etcd key.
+//
+// Keys are built as "<keyPrefix>/<namespace>/<type>/<hashedID>" by etcdKeyFromPointer, and the
+// hashed ID is hex-encoded, so it never contains a slash: cutting at the last slash yields exactly
+// the prefix etcdKeyPrefixFromKind would produce for the same kind.
+func etcdKeyPrefixFromKey(key string) string {
+	idx := strings.LastIndexByte(key, '/')
+	if idx < 0 {
+		return key
+	}
+
+	return key[:idx+1]
+}
+
 func sha256hex(input []byte) string {
 	hash := sha256.Sum256(input)
 

pkg/state/impl/etcd/options.go

@@ -24,10 +24,11 @@ type LimiterFunc func(ctx context.Context, eventType state.EventType, resourceTy
 
 // StateOptions configure etcd.State.
 type StateOptions struct {
-	observer  ObserverFunc
-	limiter   LimiterFunc
-	keyPrefix string
-	salt      []byte
+	observer    ObserverFunc
+	limiter     LimiterFunc
+	keyPrefix   string
+	salt        []byte
+	sharedWatch bool
 }
 
 // StateOption applies settings to StateOptions.
@@ -79,3 +80,33 @@ func WithLimiter(fn LimiterFunc) StateOption {
 		options.limiter = fn
 	}
 }
+
+// WithSharedWatch makes all watches created by the State share a single etcd watcher established
+// over the whole key prefix, with a single goroutine dispatching the events to the subscribers.
+//
+// Without this option every Watch/WatchKind/WatchKindAggregated call establishes its own etcd
+// watcher, and the etcd client library delivers each watcher's responses on its own goroutine, so
+// there is no ordering relationship between the streams: an event for revision N might be
+// delivered after an event for revision N+1 which belongs to a different resource type. Consumers
+// merging several watches into one stream - the COSI controller runtime being one - observe that
+// as resources appearing out of causal order.
+//
+// With this option, events delivered to the same destination channel are in etcd revision order,
+// no matter which watch produced them. The guarantee is per channel rather than global: consumers
+// routinely read several watch channels in a fixed order, and a dispatcher ordering across all of
+// them would have to block on one channel while the consumer waits on another. For the same
+// reason a slow consumer only holds up the watches writing to the channel it is reading, not every
+// watch of the State.
+//
+// The cost is that the shared watcher receives the events of every resource type under the key
+// prefix, including those nobody is watching, in exchange for establishing one etcd watcher
+// instead of one per watch call.
+//
+// Watches started with state.WithStartFromBookmark are not served by the shared watcher: it is
+// already positioned past that revision and cannot replay history. Such watches fall back to a
+// dedicated etcd watcher and are not ordered against the rest.
+func WithSharedWatch() StateOption {
+	return func(options *StateOptions) {
+		options.sharedWatch = true
+	}
+}

pkg/state/impl/etcd/state_test.go

@@ -12,6 +12,8 @@ import (
 	"github.com/cosi-project/runtime/pkg/state/conformance"
 	"github.com/stretchr/testify/suite"
 	"go.uber.org/goleak"
+
+	"github.com/cosi-project/state-etcd/pkg/state/impl/etcd"
 )
 
 func TestEtcdConformance(t *testing.T) {
@@ -24,3 +26,16 @@ func TestEtcdConformance(t *testing.T) {
 		})
 	})
 }
+
+// TestEtcdConformanceSharedWatch runs the same conformance suite with all watches served by the
+// shared watcher, as the two modes take completely different code paths.
+func TestEtcdConformanceSharedWatch(t *testing.T) {
+	defer goleak.VerifyNone(t, goleak.IgnoreCurrent())
+
+	withEtcd(t, func(s state.State) {
+		suite.Run(t, &conformance.StateSuite{
+			State:      s,
+			Namespaces: []resource.Namespace{"default", "controller", "system", "runtime"},
+		})
+	}, etcd.WithSharedWatch())
+}