clientv3: translate Snapshot API gRPC status error

To be consistent with other APIs.

Signed-off-by: Gyuho Lee <gyuhox@gmail.com>

clientv3/integration/maintenance_test.go

@@ -15,11 +15,20 @@
 package integration
 
 import (
+	"bytes"
 	"context"
+	"fmt"
+	"io"
+	"io/ioutil"
+	"path/filepath"
 	"testing"
+	"time"
 
 	"github.com/coreos/etcd/etcdserver/api/v3rpc/rpctypes"
 	"github.com/coreos/etcd/integration"
+	"github.com/coreos/etcd/lease"
+	"github.com/coreos/etcd/mvcc"
+	"github.com/coreos/etcd/mvcc/backend"
 	"github.com/coreos/etcd/pkg/testutil"
 )
 
@@ -84,3 +93,100 @@ func TestMaintenanceMoveLeader(t *testing.T) {
 		t.Fatalf("new leader expected %d, got %d", target, lead)
 	}
 }
+
+// TestMaintenanceSnapshotError ensures that context cancel/timeout
+// before snapshot reading returns corresponding context errors.
+func TestMaintenanceSnapshotError(t *testing.T) {
+	defer testutil.AfterTest(t)
+
+	clus := integration.NewClusterV3(t, &integration.ClusterConfig{Size: 1})
+	defer clus.Terminate(t)
+
+	// reading snapshot with canceled context should error out
+	ctx, cancel := context.WithCancel(context.Background())
+	rc1, err := clus.RandClient().Snapshot(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer rc1.Close()
+
+	cancel()
+	_, err = io.Copy(ioutil.Discard, rc1)
+	if err != context.Canceled {
+		t.Errorf("expected %v, got %v", context.Canceled, err)
+	}
+
+	// reading snapshot with deadline exceeded should error out
+	ctx, cancel = context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+	rc2, err := clus.RandClient().Snapshot(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer rc2.Close()
+
+	time.Sleep(2 * time.Second)
+
+	_, err = io.Copy(ioutil.Discard, rc2)
+	if err != nil && err != context.DeadlineExceeded {
+		t.Errorf("expected %v, got %v", context.DeadlineExceeded, err)
+	}
+}
+
+// TestMaintenanceSnapshotErrorInflight ensures that inflight context cancel/timeout
+// fails snapshot reading with corresponding context errors.
+func TestMaintenanceSnapshotErrorInflight(t *testing.T) {
+	defer testutil.AfterTest(t)
+
+	clus := integration.NewClusterV3(t, &integration.ClusterConfig{Size: 1})
+	defer clus.Terminate(t)
+
+	// take about 1-second to read snapshot
+	clus.Members[0].Stop(t)
+	dpath := filepath.Join(clus.Members[0].DataDir, "member", "snap", "db")
+	b := backend.NewDefaultBackend(dpath)
+	s := mvcc.NewStore(b, &lease.FakeLessor{}, nil)
+	rev := 100000
+	for i := 2; i <= rev; i++ {
+		s.Put([]byte(fmt.Sprintf("%10d", i)), bytes.Repeat([]byte("a"), 1024), lease.NoLease)
+	}
+	s.Close()
+	b.Close()
+	clus.Members[0].Restart(t)
+
+	// reading snapshot with canceled context should error out
+	ctx, cancel := context.WithCancel(context.Background())
+	rc1, err := clus.RandClient().Snapshot(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer rc1.Close()
+
+	donec := make(chan struct{})
+	go func() {
+		time.Sleep(300 * time.Millisecond)
+		cancel()
+		close(donec)
+	}()
+	_, err = io.Copy(ioutil.Discard, rc1)
+	if err != nil && err != context.Canceled {
+		t.Errorf("expected %v, got %v", context.Canceled, err)
+	}
+	<-donec
+
+	// reading snapshot with deadline exceeded should error out
+	ctx, cancel = context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+	rc2, err := clus.RandClient().Snapshot(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer rc2.Close()
+
+	// 300ms left and expect timeout while snapshot reading is in progress
+	time.Sleep(700 * time.Millisecond)
+	_, err = io.Copy(ioutil.Discard, rc2)
+	if err != nil && err != context.DeadlineExceeded {
+		t.Errorf("expected %v, got %v", context.DeadlineExceeded, err)
+	}
+}

clientv3/maintenance.go

@@ -196,7 +196,17 @@ func (m *maintenance) Snapshot(ctx context.Context) (io.ReadCloser, error) {
 		}
 		pw.Close()
 	}()
-	return pr, nil
+	return &snapshotReadCloser{ctx: ctx, ReadCloser: pr}, nil
+}
+
+type snapshotReadCloser struct {
+	ctx context.Context
+	io.ReadCloser
+}
+
+func (rc *snapshotReadCloser) Read(p []byte) (n int, err error) {
+	n, err = rc.ReadCloser.Read(p)
+	return n, toErr(rc.ctx, err)
 }
 
 func (m *maintenance) MoveLeader(ctx context.Context, transfereeID uint64) (*MoveLeaderResponse, error) {