Draft implementation of streaming range request

Documentation/dev-guide/apispec/swagger/rpc.swagger.json

@@ -2706,6 +2706,16 @@
         }
       }
     },
+    "etcdserverpbRangeStreamResponse": {
+      "type": "object",
+      "properties": {
+        "range_response": {
+          "$ref": "#/definitions/etcdserverpbRangeResponse",
+          "description": "range_response is a partial response for the KV.RangeStream RPC.\nThe result of proto.Merge() applied on all responses should result in the\nsame RangeResponse response as to the Range() method."
+        }
+      },
+      "description": "RangeStreamResponse is the response for the RangeStream RPC.\nThis message is just a wrapper around RangeResponse but there may be a need\nin the future to add streaming specific fields (for progress status, error\npropagation, etc.)."
+    },
     "etcdserverpbRequestOp": {
       "type": "object",
       "properties": {

api/etcdserverpb/rpc.proto

@@ -43,6 +43,10 @@ service KV {
     };
   }
 
+  // RangeStream gets the keys in the range from the key-value store.
+  rpc RangeStream(RangeRequest) returns (stream RangeStreamResponse) {
+  }
+
   // Put puts the given key into the key-value store.
   // A put request increments the revision of the key-value store
   // and generates one event in the event history.
@@ -498,6 +502,18 @@ message RangeRequest {
   int64 max_create_revision = 13 [(versionpb.etcd_version_field)="3.1"];
 }
 
+// RangeStreamResponse is the response for the RangeStream RPC.
+// This message is just a wrapper around RangeResponse but there may be a need
+// in the future to add streaming specific fields (for progress status, error
+// propagation, etc.).
+message RangeStreamResponse {
+  option (versionpb.etcd_version_msg) = "3.7";
+  // range_response is a partial response for the KV.RangeStream RPC.
+  // The result of proto.Merge() applied on all responses should result in the
+  // same RangeResponse response as to the Range() method.
+  RangeResponse range_response = 1;
+}
+
 message RangeResponse {
   option (versionpb.etcd_version_msg) = "3.0";
 

client/v3/kv.go

@@ -16,6 +16,8 @@ package clientv3
 
 import (
 	"context"
+	"errors"
+	"io"
 
 	"google.golang.org/grpc"
 
@@ -24,11 +26,12 @@ import (
 )
 
 type (
-	CompactResponse pb.CompactionResponse
-	PutResponse     pb.PutResponse
-	GetResponse     pb.RangeResponse
-	DeleteResponse  pb.DeleteRangeResponse
-	TxnResponse     pb.TxnResponse
+	CompactResponse   pb.CompactionResponse
+	PutResponse       pb.PutResponse
+	GetResponse       pb.RangeResponse
+	GetStreamResponse <-chan MaybeRangeStreamResponse
+	DeleteResponse    pb.DeleteRangeResponse
+	TxnResponse       pb.TxnResponse
 )
 
 type KV interface {
@@ -48,6 +51,16 @@ type KV interface {
 	// When passed WithSort(), the keys will be sorted.
 	Get(ctx context.Context, key string, opts ...OpOption) (*GetResponse, error)
 
+	// RangeStream retrieves keys.
+	// By default, will return the value for "key", if any.
+	// When passed WithRange(end), will return the keys in the range [key, end).
+	// When passed WithFromKey(), returns keys greater than or equal to key.
+	// When passed WithRev(rev) with rev > 0, retrieves keys at the given revision;
+	// if the required revision is compacted, the request will fail with ErrCompacted .
+	// When passed WithLimit(limit), the number of returned keys is bounded by limit.
+	// When passed WithSort(), the keys will be sorted.
+	GetStream(ctx context.Context, key string, opts ...OpOption) (GetStreamResponse, error)
+
 	// Delete deletes a key, or optionally using WithRange(end), [key, end).
 	Delete(ctx context.Context, key string, opts ...OpOption) (*DeleteResponse, error)
 
@@ -65,6 +78,78 @@ type KV interface {
 	Txn(ctx context.Context) Txn
 }
 
+type MaybeRangeStreamResponse struct {
+	*pb.RangeStreamResponse
+	Err error
+}
+
+func GetStreamToGetResponse(stream GetStreamResponse) (*GetResponse, error) {
+	resp := &GetResponse{Header: &pb.ResponseHeader{}}
+	for r := range stream {
+		if r.Err != nil {
+			return nil, r.Err
+		}
+		r := r.RangeStreamResponse.RangeResponse
+		if r.Header != nil {
+			if r.Header.ClusterId != 0 {
+				resp.Header.ClusterId = r.Header.ClusterId
+			}
+			if r.Header.MemberId != 0 {
+				resp.Header.MemberId = r.Header.MemberId
+			}
+			if r.Header.RaftTerm != 0 {
+				resp.Header.RaftTerm = r.Header.RaftTerm
+			}
+			if r.Header.Revision != 0 {
+				resp.Header.Revision = r.Header.Revision
+			}
+		}
+		if r.Count != 0 {
+			resp.Count = r.Count
+		}
+		if r.More {
+			resp.More = true
+		}
+		resp.Kvs = append(resp.Kvs, r.Kvs...)
+	}
+	return resp, nil
+}
+
+func RangeStreamToRangeResponse(c pb.KV_RangeStreamClient) (*pb.RangeResponse, error) {
+	resp := &pb.RangeResponse{Header: &pb.ResponseHeader{}}
+
+	for {
+		r, err := c.Recv()
+		if err != nil {
+			if !errors.Is(err, io.EOF) {
+				return nil, err
+			}
+			return resp, nil
+		}
+		if r.RangeResponse.Header != nil {
+			if r.RangeResponse.Header.ClusterId != 0 {
+				resp.Header.ClusterId = r.RangeResponse.Header.ClusterId
+			}
+			if r.RangeResponse.Header.MemberId != 0 {
+				resp.Header.MemberId = r.RangeResponse.Header.MemberId
+			}
+			if r.RangeResponse.Header.RaftTerm != 0 {
+				resp.Header.RaftTerm = r.RangeResponse.Header.RaftTerm
+			}
+			if r.RangeResponse.Header.Revision != 0 {
+				resp.Header.Revision = r.RangeResponse.Header.Revision
+			}
+		}
+		if r.RangeResponse.Count != 0 {
+			resp.Count = r.RangeResponse.Count
+		}
+		if r.RangeResponse.More {
+			resp.More = true
+		}
+		resp.Kvs = append(resp.Kvs, r.RangeResponse.Kvs...)
+	}
+}
+
 type OpResponse struct {
 	put *PutResponse
 	get *GetResponse
@@ -124,6 +209,33 @@ func (kv *kv) Get(ctx context.Context, key string, opts ...OpOption) (*GetRespon
 	return r.get, ContextError(ctx, err)
 }
 
+func (kv *kv) GetStream(ctx context.Context, key string, opts ...OpOption) (GetStreamResponse, error) {
+	op := OpGet(key, opts...)
+	c, err := kv.remote.RangeStream(ctx, op.toRangeRequest(), kv.callOpts...)
+	if err != nil {
+		return nil, ContextError(ctx, err)
+	}
+	respCh := make(chan MaybeRangeStreamResponse, 1)
+	go func() {
+		for {
+			resp, err := c.Recv()
+			if err != nil {
+				if !errors.Is(err, io.EOF) {
+					respCh <- MaybeRangeStreamResponse{
+						Err: err,
+					}
+				}
+				close(respCh)
+				return
+			}
+			respCh <- MaybeRangeStreamResponse{
+				RangeStreamResponse: resp,
+			}
+		}
+	}()
+	return respCh, nil
+}
+
 func (kv *kv) Delete(ctx context.Context, key string, opts ...OpOption) (*DeleteResponse, error) {
 	r, err := kv.Do(ctx, OpDelete(key, opts...))
 	return r.del, ContextError(ctx, err)

client/v3/leasing/kv.go

@@ -91,6 +91,10 @@ func (lkv *leasingKV) Put(ctx context.Context, key, val string, opts ...v3.OpOpt
 	return lkv.put(ctx, v3.OpPut(key, val, opts...))
 }
 
+func (lkv *leasingKV) GetStream(ctx context.Context, key string, opts ...v3.OpOption) (v3.GetStreamResponse, error) {
+	panic("unimplemented")
+}
+
 func (lkv *leasingKV) Delete(ctx context.Context, key string, opts ...v3.OpOption) (*v3.DeleteResponse, error) {
 	return lkv.delete(ctx, v3.OpDelete(key, opts...))
 }

client/v3/namespace/kv.go

@@ -64,6 +64,10 @@ func (kv *kvPrefix) Get(ctx context.Context, key string, opts ...clientv3.OpOpti
 	return get, nil
 }
 
+func (kv *kvPrefix) GetStream(ctx context.Context, key string, opts ...clientv3.OpOption) (clientv3.GetStreamResponse, error) {
+	panic("unimplemented")
+}
+
 func (kv *kvPrefix) Delete(ctx context.Context, key string, opts ...clientv3.OpOption) (*clientv3.DeleteResponse, error) {
 	if len(key) == 0 && !(clientv3.IsOptsWithFromKey(opts) || clientv3.IsOptsWithPrefix(opts)) {
 		return nil, rpctypes.ErrEmptyKey

client/v3/retry.go

@@ -108,6 +108,10 @@ func (rkv *retryKVClient) Range(ctx context.Context, in *pb.RangeRequest, opts .
 	return rkv.kc.Range(ctx, in, append(opts, withRepeatablePolicy())...)
 }
 
+func (rkv *retryKVClient) RangeStream(ctx context.Context, in *pb.RangeRequest, opts ...grpc.CallOption) (resp pb.KV_RangeStreamClient, err error) {
+	return rkv.kc.RangeStream(ctx, in, opts...)
+}
+
 func (rkv *retryKVClient) Put(ctx context.Context, in *pb.PutRequest, opts ...grpc.CallOption) (resp *pb.PutResponse, err error) {
 	return rkv.kc.Put(ctx, in, opts...)
 }

server/etcdserver/api/v3rpc/key.go

@@ -52,6 +52,29 @@ func (s *kvServer) Range(ctx context.Context, r *pb.RangeRequest) (*pb.RangeResp
 	return resp, nil
 }
 
+func (s *kvServer) RangeStream(r *pb.RangeRequest, rs pb.KV_RangeStreamServer) error {
+	if err := checkRangeRequest(r); err != nil {
+		return err
+	}
+	resp := &pb.RangeStreamResponse{
+		RangeResponse: &pb.RangeResponse{
+			Header: &pb.ResponseHeader{},
+		},
+	}
+	s.hdr.fill(resp.RangeResponse.Header)
+	// Revision will be send later
+	resp.RangeResponse.Header.Revision = 0
+	err := rs.Send(resp)
+	if err != nil {
+		return togRPCError(err)
+	}
+	err = s.kv.RangeStream(r, rs)
+	if err != nil {
+		return togRPCError(err)
+	}
+	return nil
+}
+
 func (s *kvServer) Put(ctx context.Context, r *pb.PutRequest) (*pb.PutResponse, error) {
 	if err := checkPutRequest(r); err != nil {
 		return nil, err

server/etcdserver/txn/range.go

@@ -29,7 +29,7 @@ import (
 )
 
 func Range(ctx context.Context, lg *zap.Logger, kv mvcc.KV, r *pb.RangeRequest) (resp *pb.RangeResponse, trace *traceutil.Trace, err error) {
-	ctx, trace = traceutil.EnsureTrace(ctx, lg, "range")
+	// ctx, trace = traceutil.EnsureTrace(ctx, lg, "range")
 	defer func(start time.Time) {
 		success := err == nil
 		RangeSecObserve(success, time.Since(start))
@@ -41,7 +41,7 @@ func Range(ctx context.Context, lg *zap.Logger, kv mvcc.KV, r *pb.RangeRequest)
 }
 
 func executeRange(ctx context.Context, lg *zap.Logger, txnRead mvcc.TxnRead, r *pb.RangeRequest) (*pb.RangeResponse, error) {
-	trace := traceutil.Get(ctx)
+	// trace := traceutil.Get(ctx)
 
 	limit := rangeLimit(r)
 	ro := mvcc.RangeOptions{
@@ -57,10 +57,10 @@ func executeRange(ctx context.Context, lg *zap.Logger, txnRead mvcc.TxnRead, r *
 
 	filterRangeResults(rr, r)
 	sortRangeResults(rr, r, lg)
-	trace.Step("filter and sort the key-value pairs")
+	// trace.Step("filter and sort the key-value pairs")
 
 	resp := asembleRangeResponse(rr, r)
-	trace.Step("assemble the response")
+	// trace.Step("assemble the response")
 
 	return resp, nil
 }
@@ -143,6 +143,7 @@ func asembleRangeResponse(rr *mvcc.RangeResult, r *pb.RangeRequest) *pb.RangeRes
 		rr.KVs = rr.KVs[:r.Limit]
 		resp.More = true
 	}
+	// trace.Step("filter and sort the key-value pairs")
 	resp.Header.Revision = rr.Rev
 	resp.Count = int64(rr.Count)
 	resp.Kvs = make([]*mvccpb.KeyValue, len(rr.KVs))
@@ -152,6 +153,7 @@ func asembleRangeResponse(rr *mvcc.RangeResult, r *pb.RangeRequest) *pb.RangeRes
 		}
 		resp.Kvs[i] = &rr.KVs[i]
 	}
+	// trace.Step("assemble the response")
 	return resp
 }