예제 #1
0
func TestRetryReadRows(t *testing.T) {
	ctx := context.Background()

	// Intercept requests and delegate to an interceptor defined by the test case
	errCount := 0
	var f func(grpc.ServerStream) error
	errInjector := func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
		if strings.HasSuffix(info.FullMethod, "ReadRows") {
			return f(ss)
		}
		return handler(ctx, ss)
	}

	tbl, cleanup, err := setupFakeServer(grpc.StreamInterceptor(errInjector))
	defer cleanup()
	if err != nil {
		t.Fatalf("fake server setup: %v", err)
	}

	errCount = 0
	// Test overall request failure and retries
	f = func(ss grpc.ServerStream) error {
		var err error
		req := new(btpb.ReadRowsRequest)
		ss.RecvMsg(req)
		switch errCount {
		case 0:
			// Retryable request failure
			err = grpc.Errorf(codes.Unavailable, "")
		case 1:
			// Write two rows then error
			writeReadRowsResponse(ss, "a", "b")
			err = grpc.Errorf(codes.Unavailable, "")
		case 2:
			// Retryable request failure
			if want, got := "b\x00", string(req.Rows.RowRanges[0].GetStartKeyClosed()); want != got {
				t.Errorf("2 range retries: got %q, want %q", got, want)
			}
			err = grpc.Errorf(codes.Unavailable, "")
		case 3:
			// Write two more rows
			writeReadRowsResponse(ss, "c", "d")
			err = nil
		}
		errCount++
		return err
	}

	var got []string
	tbl.ReadRows(ctx, NewRange("a", "z"), func(r Row) bool {
		got = append(got, r.Key())
		return true
	})
	want := []string{"a", "b", "c", "d"}
	if !reflect.DeepEqual(got, want) {
		t.Errorf("retry range integration: got %v, want %v", got, want)
	}
}
예제 #2
0
func Server(s *etcdserver.EtcdServer, tls *tls.Config) *grpc.Server {
	var opts []grpc.ServerOption
	opts = append(opts, grpc.CustomCodec(&codec{}))
	if tls != nil {
		opts = append(opts, grpc.Creds(credentials.NewTLS(tls)))
	}
	opts = append(opts, grpc.UnaryInterceptor(newUnaryInterceptor(s)))
	opts = append(opts, grpc.StreamInterceptor(newStreamInterceptor(s)))

	grpcServer := grpc.NewServer(opts...)
	pb.RegisterKVServer(grpcServer, NewQuotaKVServer(s))
	pb.RegisterWatchServer(grpcServer, NewWatchServer(s))
	pb.RegisterLeaseServer(grpcServer, NewQuotaLeaseServer(s))
	pb.RegisterClusterServer(grpcServer, NewClusterServer(s))
	pb.RegisterAuthServer(grpcServer, NewAuthServer(s))
	pb.RegisterMaintenanceServer(grpcServer, NewMaintenanceServer(s))

	return grpcServer
}
예제 #3
0
func (c *Component) ServerOptions() []grpc.ServerOption {
	unary := func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
		var peerAddr string
		peer, ok := peer.FromContext(ctx)
		if ok {
			peerAddr = peer.Addr.String()
		}
		var peerID string
		meta, ok := metadata.FromContext(ctx)
		if ok {
			id, ok := meta["id"]
			if ok && len(id) > 0 {
				peerID = id[0]
			}
		}
		logCtx := c.Ctx.WithFields(log.Fields{
			"CallerID": peerID,
			"CallerIP": peerAddr,
			"Method":   info.FullMethod,
		})
		t := time.Now()
		iface, err := handler(ctx, req)
		err = errors.BuildGRPCError(err)
		logCtx = logCtx.WithField("Duration", time.Now().Sub(t))
		if grpc.Code(err) == codes.OK || grpc.Code(err) == codes.Canceled {
			logCtx.Debug("Handled request")
		} else {
			logCtx.WithField("ErrCode", grpc.Code(err)).WithError(err).Debug("Handled request with error")
		}
		return iface, err
	}

	stream := func(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
		var peerAddr string
		peer, ok := peer.FromContext(stream.Context())
		if ok {
			peerAddr = peer.Addr.String()
		}
		var peerID string
		meta, ok := metadata.FromContext(stream.Context())
		if ok {
			id, ok := meta["id"]
			if ok && len(id) > 0 {
				peerID = id[0]
			}
		}
		logCtx := c.Ctx.WithFields(log.Fields{
			"CallerID": peerID,
			"CallerIP": peerAddr,
			"Method":   info.FullMethod,
		})
		t := time.Now()
		logCtx.Debug("Start stream")
		err := handler(srv, stream)
		err = errors.BuildGRPCError(err)
		logCtx = logCtx.WithField("Duration", time.Now().Sub(t))
		if grpc.Code(err) == codes.OK || grpc.Code(err) == codes.Canceled {
			logCtx.Debug("End stream")
		} else {
			logCtx.WithField("ErrCode", grpc.Code(err)).WithError(err).Debug("End stream with error")
		}
		return err
	}

	opts := []grpc.ServerOption{
		grpc.UnaryInterceptor(grpc_middleware.ChainUnaryServer(unary)),
		grpc.StreamInterceptor(grpc_middleware.ChainStreamServer(stream)),
	}

	if c.tlsConfig != nil {
		opts = append(opts, grpc.Creds(credentials.NewTLS(c.tlsConfig)))
	}

	return opts
}
예제 #4
0
func startGRPCProxy(cmd *cobra.Command, args []string) {
	l, err := net.Listen("tcp", grpcProxyListenAddr)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
	if l, err = transport.NewKeepAliveListener(l, "tcp", nil); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
	plog.Infof("listening for grpc-proxy client requests on %s", grpcProxyListenAddr)
	defer func() {
		l.Close()
		plog.Infof("stopping listening for grpc-proxy client requests on %s", grpcProxyListenAddr)
	}()
	m := cmux.New(l)

	cfg, err := newClientCfg()
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	client, err := clientv3.New(*cfg)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	kvp := grpcproxy.NewKvProxy(client)
	watchp := grpcproxy.NewWatchProxy(client)
	clusterp := grpcproxy.NewClusterProxy(client)
	leasep := grpcproxy.NewLeaseProxy(client)
	mainp := grpcproxy.NewMaintenanceProxy(client)
	authp := grpcproxy.NewAuthProxy(client)

	server := grpc.NewServer(
		grpc.StreamInterceptor(grpc_prometheus.StreamServerInterceptor),
		grpc.UnaryInterceptor(grpc_prometheus.UnaryServerInterceptor),
	)
	pb.RegisterKVServer(server, kvp)
	pb.RegisterWatchServer(server, watchp)
	pb.RegisterClusterServer(server, clusterp)
	pb.RegisterLeaseServer(server, leasep)
	pb.RegisterMaintenanceServer(server, mainp)
	pb.RegisterAuthServer(server, authp)

	errc := make(chan error)

	grpcl := m.Match(cmux.HTTP2())
	go func() { errc <- server.Serve(grpcl) }()

	httpmux := http.NewServeMux()
	httpmux.HandleFunc("/", http.NotFound)
	httpmux.Handle("/metrics", prometheus.Handler())
	srvhttp := &http.Server{
		Handler: httpmux,
	}

	var httpl net.Listener
	if cfg.TLS != nil {
		srvhttp.TLSConfig = cfg.TLS
		httpl = tls.NewListener(m.Match(cmux.Any()), cfg.TLS)
	} else {
		httpl = m.Match(cmux.HTTP1())
	}
	go func() { errc <- srvhttp.Serve(httpl) }()

	go func() { errc <- m.Serve() }()

	fmt.Fprintln(os.Stderr, <-errc)
	os.Exit(1)
}
예제 #5
0
func TestRetryApplyBulk(t *testing.T) {
	ctx := context.Background()

	// Intercept requests and delegate to an interceptor defined by the test case
	errCount := 0
	var f func(grpc.ServerStream) error
	errInjector := func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
		if strings.HasSuffix(info.FullMethod, "MutateRows") {
			return f(ss)
		}
		return handler(ctx, ss)
	}

	tbl, cleanup, err := setupFakeServer(grpc.StreamInterceptor(errInjector))
	defer cleanup()
	if err != nil {
		t.Fatalf("fake server setup: %v", err)
	}

	errCount = 0
	// Test overall request failure and retries
	f = func(ss grpc.ServerStream) error {
		if errCount < 3 {
			errCount++
			return grpc.Errorf(codes.Aborted, "")
		}
		return nil
	}
	mut := NewMutation()
	mut.Set("cf", "col", 1, []byte{})
	errors, err := tbl.ApplyBulk(ctx, []string{"row2"}, []*Mutation{mut})
	if errors != nil || err != nil {
		t.Errorf("bulk with request failure: got: %v, %v, want: nil", errors, err)
	}

	// Test failures and retries in one request
	errCount = 0
	m1 := NewMutation()
	m1.Set("cf", "col", 1, []byte{})
	m2 := NewMutation()
	m2.Set("cf", "col2", 1, []byte{})
	m3 := NewMutation()
	m3.Set("cf", "col3", 1, []byte{})
	f = func(ss grpc.ServerStream) error {
		var err error
		req := new(btpb.MutateRowsRequest)
		ss.RecvMsg(req)
		switch errCount {
		case 0:
			// Retryable request failure
			err = grpc.Errorf(codes.Unavailable, "")
		case 1:
			// Two mutations fail
			writeMutateRowsResponse(ss, codes.Unavailable, codes.OK, codes.Aborted)
			err = nil
		case 2:
			// Two failures were retried. One will succeed.
			if want, got := 2, len(req.Entries); want != got {
				t.Errorf("2 bulk retries, got: %d, want %d", got, want)
			}
			writeMutateRowsResponse(ss, codes.OK, codes.Aborted)
			err = nil
		case 3:
			// One failure was retried and will succeed.
			if want, got := 1, len(req.Entries); want != got {
				t.Errorf("1 bulk retry, got: %d, want %d", got, want)
			}
			writeMutateRowsResponse(ss, codes.OK)
			err = nil
		}
		errCount++
		return err
	}
	errors, err = tbl.ApplyBulk(ctx, []string{"row1", "row2", "row3"}, []*Mutation{m1, m2, m3})
	if errors != nil || err != nil {
		t.Errorf("bulk with retries: got: %v, %v, want: nil", errors, err)
	}

	// Test unretryable errors
	niMut := NewMutation()
	niMut.Set("cf", "col", ServerTime, []byte{}) // Non-idempotent
	errCount = 0
	f = func(ss grpc.ServerStream) error {
		var err error
		req := new(btpb.MutateRowsRequest)
		ss.RecvMsg(req)
		switch errCount {
		case 0:
			// Give non-idempotent mutation a retryable error code.
			// Nothing should be retried.
			writeMutateRowsResponse(ss, codes.FailedPrecondition, codes.Aborted)
			err = nil
		case 1:
			t.Errorf("unretryable errors: got one retry, want no retries")
		}
		errCount++
		return err
	}
	errors, err = tbl.ApplyBulk(ctx, []string{"row1", "row2"}, []*Mutation{m1, niMut})
	if err != nil {
		t.Errorf("unretryable errors: request failed %v")
	}
	want := []error{
		grpc.Errorf(codes.FailedPrecondition, ""),
		grpc.Errorf(codes.Aborted, ""),
	}
	if !reflect.DeepEqual(want, errors) {
		t.Errorf("unretryable errors: got: %v, want: %v", errors, want)
	}

	// Test individual errors and a deadline exceeded
	f = func(ss grpc.ServerStream) error {
		writeMutateRowsResponse(ss, codes.FailedPrecondition, codes.OK, codes.Aborted)
		return nil
	}
	ctx, _ = context.WithTimeout(ctx, 100*time.Millisecond)
	errors, err = tbl.ApplyBulk(ctx, []string{"row1", "row2", "row3"}, []*Mutation{m1, m2, m3})
	wantErr := context.DeadlineExceeded
	if wantErr != err {
		t.Errorf("deadline exceeded error: got: %v, want: %v", err, wantErr)
	}
	if errors != nil {
		t.Errorf("deadline exceeded errors: got: %v, want: nil", err)
	}
}