CLI implementation for get sequence updates (#690)

(stacked on top of #689 )

Author: merlimat

cmd/client/cmd.go

@@ -17,6 +17,8 @@ package client
 import (
 	"fmt"
 
+	"github.com/streamnative/oxia/cmd/client/sequenceupdates"
+
 	"github.com/spf13/cobra"
 
 	"github.com/streamnative/oxia/cmd/client/common"
@@ -52,4 +54,5 @@ func init() {
 	Cmd.AddCommand(rangescan.Cmd)
 	Cmd.AddCommand(deleterange.Cmd)
 	Cmd.AddCommand(notifications.Cmd)
+	Cmd.AddCommand(sequenceupdates.Cmd)
 }

cmd/client/common/client_mock.go

@@ -81,3 +81,7 @@ func (m *MockClient) RangeScan(_ context.Context, minKeyInclusive string, maxKey
 func (*MockClient) GetNotifications() (oxia.Notifications, error) {
 	return nil, errors.New("not implemented in mock")
 }
+
+func (*MockClient) GetSequenceUpdates(context.Context, string, ...oxia.GetSequenceUpdatesOption) (<-chan string, error) {
+	return nil, errors.New("not implemented in mock")
+}

cmd/client/sequenceupdates/cmd.go

@@ -0,0 +1,71 @@
+// Copyright 2023 StreamNative, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package sequenceupdates
+
+import (
+	"context"
+	"log/slog"
+
+	"github.com/spf13/cobra"
+
+	"github.com/streamnative/oxia/cmd/client/common"
+	"github.com/streamnative/oxia/oxia"
+)
+
+var Config = flags{}
+
+type flags struct {
+	key          string
+	partitionKey string
+}
+
+func (flags *flags) Reset() {
+	flags.key = ""
+	flags.partitionKey = ""
+}
+
+func init() {
+	Cmd.Flags().StringVarP(&Config.partitionKey, "partition-key", "p", "", "Partition Key to be used in override the shard routing")
+	_ = Cmd.MarkFlagRequired("partition-key")
+}
+
+var Cmd = &cobra.Command{
+	Use:   "sequence-updates [flags] KEY --partition-key PARTITION_KEY",
+	Short: "Get key sequences updates",
+	Long:  `Follow the updates on a sequence of keys`,
+	Args:  cobra.ExactArgs(1),
+	RunE:  exec,
+}
+
+func exec(_ *cobra.Command, args []string) error {
+	client, err := common.Config.NewClient()
+	if err != nil {
+		return err
+	}
+
+	defer client.Close()
+
+	Config.key = args[0]
+	updates, err := client.GetSequenceUpdates(context.Background(), Config.key, oxia.PartitionKey(Config.partitionKey))
+	if err != nil {
+		return err
+	}
+
+	for key := range updates {
+		slog.Info("", slog.String("key", key))
+	}
+
+	return nil
+}

oxia/async_client_impl.go

@@ -458,6 +458,15 @@ func (c *clientImpl) RangeScan(ctx context.Context, minKeyInclusive string, maxK
 	return outCh
 }
 
+func (c *clientImpl) GetSequenceUpdates(ctx context.Context, prefixKey string, options ...GetSequenceUpdatesOption) (<-chan string, error) {
+	opts := newGetSequenceUpdatesOptions(options)
+	if opts.partitionKey == nil {
+		return nil, errors.Wrap(ErrInvalidOptions, "partitionKey is required")
+	}
+
+	return newSequenceUpdates(ctx, prefixKey, *opts.partitionKey, c.clientPool, c.shardManager), nil
+}
+
 // We do range scan on all the shards, and we need to always pick the lowest key
 // across all the shards.
 func aggregateAndSortRangeScanAcrossShards(channels []chan GetResult, outCh chan GetResult) {

oxia/client.go

@@ -103,6 +103,12 @@ type AsyncClient interface {
 	// https://github.com/streamnative/oxia/blob/main/docs/oxia-key-sorting.md
 	RangeScan(ctx context.Context, minKeyInclusive string, maxKeyExclusive string, options ...RangeScanOption) <-chan GetResult
 
+	// GetSequenceUpdates allows to subscribe to the updates happening on a sequential key
+	// The channel will report the current latest sequence for a given key.
+	// Multiple updates can be collapsed into one single event with the
+	// highest sequence.
+	GetSequenceUpdates(ctx context.Context, prefixKey string, options ...GetSequenceUpdatesOption) (<-chan string, error)
+
 	// GetNotifications creates a new subscription to receive the notifications
 	// from Oxia for any change that is applied to the database
 	GetNotifications() (Notifications, error)
@@ -164,6 +170,12 @@ type SyncClient interface {
 	// inserted with that partition key).
 	RangeScan(ctx context.Context, minKeyInclusive string, maxKeyExclusive string, options ...RangeScanOption) <-chan GetResult
 
+	// GetSequenceUpdates allows to subscribe to the updates happening on a sequential key
+	// The channel will report the current latest sequence for a given key.
+	// Multiple updates can be collapsed into one single event with the
+	// highest sequence.
+	GetSequenceUpdates(ctx context.Context, prefixKey string, options ...GetSequenceUpdatesOption) (<-chan string, error)
+
 	// GetNotifications creates a new subscription to receive the notifications
 	// from Oxia for any change that is applied to the database
 	GetNotifications() (Notifications, error)

oxia/options_base.go

@@ -22,6 +22,7 @@ type BaseOption interface {
 	DeleteRangeOption
 	ListOption
 	RangeScanOption
+	GetSequenceUpdatesOption
 }
 
 type baseOptions struct {
@@ -66,6 +67,10 @@ func (o *partitionKeyOpt) applyRangeScan(opts *rangeScanOptions) {
 	opts.partitionKey = o.partitionKey
 }
 
+func (o *partitionKeyOpt) applyGetSequenceUpdates(opts *getSequenceUpdatesOptions) {
+	opts.partitionKey = o.partitionKey
+}
+
 // PartitionKey overrides the partition routing with the specified `partitionKey` instead
 // of the regular record key.
 // Records with the same partitionKey will always be guaranteed to be co-located in the

oxia/options_get_sequence_updates.go

@@ -0,0 +1,32 @@
+// Copyright 2023 StreamNative, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package oxia
+
+type getSequenceUpdatesOptions struct {
+	baseOptions
+}
+
+// GetSequenceUpdatesOption represents an option for the [SyncClient.GetSequenceUpdates] operation.
+type GetSequenceUpdatesOption interface {
+	applyGetSequenceUpdates(opts *getSequenceUpdatesOptions)
+}
+
+func newGetSequenceUpdatesOptions(opts []GetSequenceUpdatesOption) *getSequenceUpdatesOptions {
+	getSequenceUpdatesOptions := &getSequenceUpdatesOptions{}
+	for _, opt := range opts {
+		opt.applyGetSequenceUpdates(getSequenceUpdatesOptions)
+	}
+	return getSequenceUpdatesOptions
+}

oxia/sequence_updates.go

@@ -0,0 +1,124 @@
+// Copyright 2023 StreamNative, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package oxia
+
+import (
+	"context"
+	"errors"
+	"log/slog"
+	"time"
+
+	"github.com/cenkalti/backoff/v4"
+
+	"github.com/streamnative/oxia/common"
+	"github.com/streamnative/oxia/oxia/internal"
+	"github.com/streamnative/oxia/proto"
+)
+
+type sequenceUpdates struct {
+	prefixKey    string
+	partitionKey string
+	ch           chan string
+	shardManager internal.ShardManager
+	clientPool   common.ClientPool
+
+	ctx     context.Context
+	backoff backoff.BackOff
+	log     *slog.Logger
+}
+
+func newSequenceUpdates(ctx context.Context, prefixKey string, partitionKey string,
+	clientPool common.ClientPool, shardManager internal.ShardManager) <-chan string {
+	su := &sequenceUpdates{
+		prefixKey:    prefixKey,
+		partitionKey: partitionKey,
+		ch:           make(chan string),
+		shardManager: shardManager,
+		clientPool:   clientPool,
+		ctx:          ctx,
+		backoff:      common.NewBackOffWithInitialInterval(ctx, 1*time.Second),
+		log: slog.With(
+			slog.String("component", "oxia-get-sequence-updates"),
+			slog.String("key", "key"),
+		),
+	}
+
+	go common.DoWithLabels(
+		su.ctx,
+		map[string]string{
+			"oxia":      "sequence-updates",
+			"prefixKey": prefixKey,
+		},
+		su.getSequenceUpdatesWithRetries,
+	)
+
+	return su.ch
+}
+
+func (su *sequenceUpdates) getSequenceUpdatesWithRetries() { //nolint:revive
+	_ = backoff.RetryNotify(su.getSequenceUpdates,
+		su.backoff, func(err error, duration time.Duration) {
+			if !errors.Is(err, context.Canceled) {
+				su.log.Error(
+					"Error while getting sequence updates",
+					slog.Any("error", err),
+					slog.Duration("retry-after", duration),
+				)
+			}
+		})
+
+	// Signal that the background go-routine is now done
+	close(su.ch)
+}
+
+func (su *sequenceUpdates) getSequenceUpdates() error {
+	shard := su.shardManager.Get(su.partitionKey)
+	leader := su.shardManager.Leader(shard)
+
+	rpc, err := su.clientPool.GetClientRpc(leader)
+	if err != nil {
+		return err
+	}
+
+	updates, err := rpc.GetSequenceUpdates(su.ctx, &proto.GetSequenceUpdatesRequest{
+		Key: su.prefixKey,
+	})
+	if err != nil {
+		if su.ctx.Err() != nil {
+			return su.ctx.Err()
+		}
+		return err
+	}
+
+	su.backoff.Reset()
+
+	for {
+		res, err2 := updates.Recv()
+		if err2 != nil {
+			return err2
+		}
+
+		if res.HighestSequenceKey == "" {
+			// Ignore first response if there are no sequences for the key
+			continue
+		}
+
+		select {
+		case su.ch <- res.HighestSequenceKey:
+		case <-su.ctx.Done():
+			return su.ctx.Err()
+		}
+	}
+}

oxia/sync_client_impl.go

@@ -133,6 +133,10 @@ func (c *syncClientImpl) RangeScan(ctx context.Context, minKeyInclusive string,
 	return c.asyncClient.RangeScan(ctx, minKeyInclusive, maxKeyExclusive, options...)
 }
 
+func (c *syncClientImpl) GetSequenceUpdates(ctx context.Context, prefixKey string, options ...GetSequenceUpdatesOption) (<-chan string, error) {
+	return c.asyncClient.GetSequenceUpdates(ctx, prefixKey, options...)
+}
+
 func (c *syncClientImpl) GetNotifications() (Notifications, error) {
 	return c.asyncClient.GetNotifications()
 }