// TestV3LeaseExists creates a lease on a random client, then sends a keepalive on another
// client to confirm it's visible to the whole cluster.
func TestV3LeaseExists(t *testing.T) {
clus := newClusterGRPC(t, &clusterConfig{size: 3})
defer clus.Terminate(t)
// create lease
lresp, err := pb.NewLeaseClient(clus.RandConn()).LeaseCreate(
context.TODO(),
&pb.LeaseCreateRequest{TTL: 30})
if err != nil {
t.Fatal(err)
}
if lresp.Error != "" {
t.Fatal(lresp.Error)
}
// confirm keepalive
lac, err := pb.NewLeaseClient(clus.RandConn()).LeaseKeepAlive(context.TODO())
if err != nil {
t.Fatal(err)
}
defer lac.CloseSend()
if err = lac.Send(&pb.LeaseKeepAliveRequest{ID: lresp.ID}); err != nil {
t.Fatal(err)
}
if _, err = lac.Recv(); err != nil {
t.Fatal(err)
}
}
// TestV3LeaseExists creates a lease on a random client, then sends a keepalive on another
// client to confirm it's visible to the whole cluster.
func TestV3LeaseExists(t *testing.T) {
defer testutil.AfterTest(t)
clus := newClusterGRPC(t, &clusterConfig{size: 3})
defer clus.Terminate(t)
// create lease
ctx0, cancel0 := context.WithCancel(context.Background())
defer cancel0()
lresp, err := pb.NewLeaseClient(clus.RandConn()).LeaseCreate(
ctx0,
&pb.LeaseCreateRequest{TTL: 30})
if err != nil {
t.Fatal(err)
}
if lresp.Error != "" {
t.Fatal(lresp.Error)
}
// confirm keepalive
ctx1, cancel1 := context.WithCancel(context.Background())
defer cancel1()
lac, err := pb.NewLeaseClient(clus.RandConn()).LeaseKeepAlive(ctx1)
if err != nil {
t.Fatal(err)
}
defer lac.CloseSend()
if err = lac.Send(&pb.LeaseKeepAliveRequest{ID: lresp.ID}); err != nil {
t.Fatal(err)
}
if _, err = lac.Recv(); err != nil {
t.Fatal(err)
}
}
func newClient(cfg *Config) (*Client, error) {
if cfg == nil {
cfg = &Config{RetryDialer: dialEndpointList}
}
var creds *credentials.TransportAuthenticator
if cfg.TLS != nil {
tlscfg, err := cfg.TLS.ClientConfig()
if err != nil {
return nil, err
}
c := credentials.NewTLS(tlscfg)
creds = &c
}
// use a temporary skeleton client to bootstrap first connection
conn, err := cfg.RetryDialer(&Client{cfg: *cfg, creds: creds})
if err != nil {
return nil, err
}
return &Client{
KV: pb.NewKVClient(conn),
Lease: pb.NewLeaseClient(conn),
Watch: pb.NewWatchClient(conn),
Cluster: pb.NewClusterClient(conn),
conn: conn,
cfg: *cfg,
creds: creds,
}, nil
}
func (l *lessor) switchRemoteAndStream(prevErr error) error {
l.mu.Lock()
conn := l.conn
l.mu.Unlock()
var (
err error
newConn *grpc.ClientConn
)
if prevErr != nil {
conn.Close()
newConn, err = l.c.retryConnection(conn, prevErr)
if err != nil {
return err
}
}
l.mu.Lock()
if newConn != nil {
l.conn = newConn
}
l.remote = pb.NewLeaseClient(l.conn)
l.mu.Unlock()
serr := l.newStream()
if serr != nil {
return serr
}
return nil
}
// TestV3LeaseKeepAlive ensures keepalive keeps the lease alive.
func TestV3LeaseKeepAlive(t *testing.T) {
testLeaseRemoveLeasedKey(t, func(clus *clusterV3, leaseID int64) error {
lc := pb.NewLeaseClient(clus.RandConn())
lreq := &pb.LeaseKeepAliveRequest{ID: leaseID}
lac, err := lc.LeaseKeepAlive(context.TODO())
if err != nil {
return err
}
defer lac.CloseSend()
// renew long enough so lease would've expired otherwise
for i := 0; i < 3; i++ {
if err = lac.Send(lreq); err != nil {
return err
}
lresp, rxerr := lac.Recv()
if rxerr != nil {
return rxerr
}
if lresp.ID != leaseID {
return fmt.Errorf("expected lease ID %v, got %v", leaseID, lresp.ID)
}
time.Sleep(time.Duration(lresp.TTL/2) * time.Second)
}
_, err = lc.LeaseRevoke(context.TODO(), &pb.LeaseRevokeRequest{ID: leaseID})
return err
})
}
// leaseRevokeCommandFunc executes the "lease create" command.
func leaseRevokeCommandFunc(cmd *cobra.Command, args []string) {
if len(args) != 1 {
ExitWithError(ExitBadArgs, fmt.Errorf("lease revoke command needs 1 argument"))
}
id, err := strconv.ParseInt(args[0], 16, 64)
if err != nil {
ExitWithError(ExitBadArgs, fmt.Errorf("bad lease ID arg (%v), expecting ID in Hex", err))
}
endpoint, err := cmd.Flags().GetString("endpoint")
if err != nil {
ExitWithError(ExitError, err)
}
conn, err := grpc.Dial(endpoint)
if err != nil {
ExitWithError(ExitBadConnection, err)
}
lease := pb.NewLeaseClient(conn)
req := &pb.LeaseRevokeRequest{ID: id}
_, err = lease.LeaseRevoke(context.Background(), req)
if err != nil {
fmt.Fprintf(os.Stderr, "failed to revoke lease (%v)\n", err)
return
}
fmt.Printf("lease %016x revoked\n", id)
}
// leaseCreateCommandFunc executes the "lease create" command.
func leaseCreateCommandFunc(cmd *cobra.Command, args []string) {
if len(args) != 1 {
ExitWithError(ExitBadArgs, fmt.Errorf("lease create command needs TTL argument."))
}
ttl, err := strconv.ParseInt(args[0], 10, 64)
if err != nil {
ExitWithError(ExitBadArgs, fmt.Errorf("bad TTL (%v)", err))
}
endpoint, err := cmd.Flags().GetString("endpoint")
if err != nil {
ExitWithError(ExitError, err)
}
conn, err := grpc.Dial(endpoint)
if err != nil {
ExitWithError(ExitBadConnection, err)
}
lease := pb.NewLeaseClient(conn)
req := &pb.LeaseCreateRequest{TTL: ttl}
resp, err := lease.LeaseCreate(context.Background(), req)
if err != nil {
fmt.Fprintf(os.Stderr, "failed to create lease (%v)\n", err)
return
}
fmt.Printf("lease %016x created with TTL(%ds)\n", resp.ID, resp.TTL)
}
// TestV3LeaseRevoke ensures a key is deleted once its lease is revoked.
func TestV3LeaseRevoke(t *testing.T) {
testLeaseRemoveLeasedKey(t, func(clus *clusterV3, leaseID int64) error {
lc := pb.NewLeaseClient(clus.RandConn())
_, err := lc.LeaseRevoke(context.TODO(), &pb.LeaseRevokeRequest{ID: leaseID})
return err
})
}
func toGRPC(c *clientv3.Client) grpcAPI {
return grpcAPI{
pb.NewClusterClient(c.ActiveConnection()),
pb.NewKVClient(c.ActiveConnection()),
pb.NewLeaseClient(c.ActiveConnection()),
pb.NewWatchClient(c.ActiveConnection()),
}
}
func toGRPC(c *clientv3.Client) grpcAPI {
if v, ok := proxies[c]; ok {
return v
}
return grpcAPI{
pb.NewClusterClient(c.ActiveConnection()),
grpcproxy.KvServerToKvClient(grpcproxy.NewKvProxy(c)),
pb.NewLeaseClient(c.ActiveConnection()),
grpcproxy.WatchServerToWatchClient(grpcproxy.NewWatchProxy(c)),
pb.NewMaintenanceClient(c.ActiveConnection()),
}
}
func newClient(conn *grpc.ClientConn, cfg *Config) *Client {
if cfg == nil {
cfg = &Config{RetryDialer: dialEndpointList}
}
return &Client{
KV: pb.NewKVClient(conn),
Lease: pb.NewLeaseClient(conn),
Watch: pb.NewWatchClient(conn),
Cluster: pb.NewClusterClient(conn),
conn: conn,
cfg: *cfg,
}
}
// TestV3LeaseCreateById ensures leases may be created by a given id.
func TestV3LeaseCreateByID(t *testing.T) {
defer testutil.AfterTest(t)
clus := newClusterGRPC(t, &clusterConfig{size: 3})
defer clus.Terminate(t)
// create fixed lease
lresp, err := pb.NewLeaseClient(clus.RandConn()).LeaseCreate(
context.TODO(),
&pb.LeaseCreateRequest{ID: 1, TTL: 1})
if err != nil {
t.Errorf("could not create lease 1 (%v)", err)
}
if lresp.ID != 1 {
t.Errorf("got id %v, wanted id %v", lresp.ID)
}
// create duplicate fixed lease
lresp, err = pb.NewLeaseClient(clus.RandConn()).LeaseCreate(
context.TODO(),
&pb.LeaseCreateRequest{ID: 1, TTL: 1})
if err != nil {
t.Error(err)
}
if lresp.ID != 0 || lresp.Error != lease.ErrLeaseExists.Error() {
t.Errorf("got id %v, wanted id 0 (%v)", lresp.ID, lresp.Error)
}
// create fresh fixed lease
lresp, err = pb.NewLeaseClient(clus.RandConn()).LeaseCreate(
context.TODO(),
&pb.LeaseCreateRequest{ID: 2, TTL: 1})
if err != nil {
t.Errorf("could not create lease 2 (%v)", err)
}
if lresp.ID != 2 {
t.Errorf("got id %v, wanted id %v", lresp.ID)
}
}
func NewLease(c *Client) Lease {
l := &lessor{
donec: make(chan struct{}),
keepAlives: make(map[LeaseID]*keepAlive),
}
f := func(conn *grpc.ClientConn) { l.remote = pb.NewLeaseClient(conn) }
l.rc = newRemoteClient(c, f)
l.stopCtx, l.stopCancel = context.WithCancel(context.Background())
go l.recvKeepAliveLoop()
return l
}
func NewLease(c *Client) Lease {
l := &lessor{
donec: make(chan struct{}),
keepAlives: make(map[LeaseID]*keepAlive),
remote: pb.NewLeaseClient(c.conn),
firstKeepAliveTimeout: c.cfg.DialTimeout + time.Second,
}
if l.firstKeepAliveTimeout == time.Second {
l.firstKeepAliveTimeout = defaultTTL
}
l.stopCtx, l.stopCancel = context.WithCancel(context.Background())
go l.recvKeepAliveLoop()
go l.deadlineLoop()
return l
}
func NewLease(c *Client) Lease {
l := &lessor{
c: c,
conn: c.ActiveConnection(),
donec: make(chan struct{}),
keepAlives: make(map[LeaseID]*keepAlive),
}
l.remote = pb.NewLeaseClient(l.conn)
l.stopCtx, l.stopCancel = context.WithCancel(context.Background())
go l.recvKeepAliveLoop()
return l
}
func toGRPC(c *clientv3.Client) grpcAPI {
pmu.Lock()
defer pmu.Unlock()
if v, ok := proxies[c]; ok {
return v
}
api := grpcAPI{
pb.NewClusterClient(c.ActiveConnection()),
grpcproxy.KvServerToKvClient(grpcproxy.NewKvProxy(c)),
pb.NewLeaseClient(c.ActiveConnection()),
grpcproxy.WatchServerToWatchClient(grpcproxy.NewWatchProxy(c)),
pb.NewMaintenanceClient(c.ActiveConnection()),
}
proxies[c] = api
return api
}
// acquireLeaseAndKey creates a new lease and creates an attached key.
func acquireLeaseAndKey(clus *clusterV3, key string) (int64, error) {
// create lease
lresp, err := pb.NewLeaseClient(clus.RandConn()).LeaseCreate(
context.TODO(),
&pb.LeaseCreateRequest{TTL: 1})
if err != nil {
return 0, err
}
if lresp.Error != "" {
return 0, fmt.Errorf(lresp.Error)
}
// attach to key
put := &pb.PutRequest{Key: []byte(key), Lease: lresp.ID}
if _, err := pb.NewKVClient(clus.RandConn()).Put(context.TODO(), put); err != nil {
return 0, err
}
return lresp.ID, nil
}
func (lc *leaseChecker) Check() error {
conn, err := grpc.Dial(lc.ls.endpoint, grpc.WithInsecure(), grpc.WithBackoffMaxDelay(1))
if err != nil {
return fmt.Errorf("%v (%s)", err, lc.ls.endpoint)
}
defer func() {
if conn != nil {
conn.Close()
}
}()
lc.kvc = pb.NewKVClient(conn)
lc.leaseClient = pb.NewLeaseClient(conn)
if err := lc.check(true, lc.ls.revokedLeases.leases); err != nil {
return err
}
if err := lc.check(false, lc.ls.aliveLeases.leases); err != nil {
return err
}
return lc.checkShortLivedLeases()
}
func NewLease(c *Client) Lease {
l := &lessor{
c: c,
conn: c.ActiveConnection(),
initedc: make(chan bool, 1),
keepAlives: make(map[lease.LeaseID]chan *LeaseKeepAliveResponse),
deadlines: make(map[lease.LeaseID]time.Time),
}
l.remote = pb.NewLeaseClient(l.conn)
l.stopCtx, l.stopCancel = context.WithCancel(context.Background())
l.initedc <- false
go l.recvKeepAliveLoop()
go l.sendKeepAliveLoop()
return l
}
func (lp *leaseProxy) LeaseKeepAlive(stream pb.Lease_LeaseKeepAliveServer) error {
conn := lp.client.ActiveConnection()
ctx, cancel := context.WithCancel(stream.Context())
lc, err := pb.NewLeaseClient(conn).LeaseKeepAlive(ctx)
if err != nil {
cancel()
return err
}
go func() {
// Cancel the context attached to lc to unblock lc.Recv when
// this routine returns on error.
defer cancel()
for {
// stream.Recv will be unblock when the loop in the parent routine
// returns on error.
rr, err := stream.Recv()
if err != nil {
return
}
err = lc.Send(rr)
if err != nil {
return
}
}
}()
for {
rr, err := lc.Recv()
if err != nil {
return err
}
err = stream.Send(rr)
if err != nil {
return err
}
}
}
func newClient(conn *grpc.ClientConn, cfg *Config) (*Client, error) {
if cfg == nil {
cfg = &Config{RetryDialer: dialEndpointList}
}
var creds *credentials.TransportAuthenticator
if cfg.TLS != nil {
tlscfg, err := cfg.TLS.ClientConfig()
if err != nil {
return nil, err
}
c := credentials.NewTLS(tlscfg)
creds = &c
}
return &Client{
KV: pb.NewKVClient(conn),
Lease: pb.NewLeaseClient(conn),
Watch: pb.NewWatchClient(conn),
Cluster: pb.NewClusterClient(conn),
conn: conn,
cfg: *cfg,
creds: creds,
}, nil
}
func (ls *leaseStresser) setupOnce() error {
if ls.aliveLeases != nil {
return nil
}
if ls.numLeases == 0 {
panic("expect numLeases to be set")
}
if ls.keysPerLease == 0 {
panic("expect keysPerLease to be set")
}
conn, err := grpc.Dial(ls.endpoint, grpc.WithInsecure())
if err != nil {
return fmt.Errorf("%v (%s)", err, ls.endpoint)
}
ls.conn = conn
ls.kvc = pb.NewKVClient(conn)
ls.lc = pb.NewLeaseClient(conn)
ls.aliveLeases = &atomicLeases{leases: make(map[int64]time.Time)}
return nil
}
// leaseKeepAliveCommandFunc executes the "lease keep-alive" command.
func leaseKeepAliveCommandFunc(cmd *cobra.Command, args []string) {
if len(args) != 1 {
ExitWithError(ExitBadArgs, fmt.Errorf("lease keep-alive command needs lease ID as argument"))
}
id, err := strconv.ParseInt(args[0], 16, 64)
if err != nil {
ExitWithError(ExitBadArgs, fmt.Errorf("bad lease ID arg (%v), expecting ID in Hex", err))
}
endpoint, err := cmd.Flags().GetString("endpoint")
if err != nil {
ExitWithError(ExitError, err)
}
conn, err := grpc.Dial(endpoint)
if err != nil {
ExitWithError(ExitBadConnection, err)
}
lease := pb.NewLeaseClient(conn)
kStream, err := lease.LeaseKeepAlive(context.TODO())
if err != nil {
ExitWithError(ExitBadConnection, err)
}
nextC := make(chan int64, 1)
go leaseKeepAliveRecvLoop(kStream, nextC)
req := &pb.LeaseKeepAliveRequest{ID: id}
for {
err := kStream.Send(req)
if err != nil {
ExitWithError(ExitError, fmt.Errorf("failed to keep-alive lease (%v)", err))
}
next := <-nextC
time.Sleep(time.Duration(next/2) * time.Second)
}
}
func (ls *leaseStresser) Stress() error {
plog.Infof("lease Stresser %v starting ...", ls.endpoint)
if err := ls.setupOnce(); err != nil {
return err
}
conn, err := grpc.Dial(ls.endpoint, grpc.WithInsecure(), grpc.WithBackoffMaxDelay(1*time.Second))
if err != nil {
return fmt.Errorf("%v (%s)", err, ls.endpoint)
}
ls.conn = conn
ls.kvc = pb.NewKVClient(conn)
ls.lc = pb.NewLeaseClient(conn)
ls.revokedLeases = &atomicLeases{leases: make(map[int64]time.Time)}
ls.shortLivedLeases = &atomicLeases{leases: make(map[int64]time.Time)}
ctx, cancel := context.WithCancel(context.Background())
ls.cancel = cancel
ls.ctx = ctx
ls.runWg.Add(1)
go ls.run()
return nil
}
func (lp *leaseProxy) LeaseRevoke(ctx context.Context, rr *pb.LeaseRevokeRequest) (*pb.LeaseRevokeResponse, error) {
conn := lp.client.ActiveConnection()
return pb.NewLeaseClient(conn).LeaseRevoke(ctx, rr)
}
func NewEtcdClient(conn *grpc.ClientConn) *EtcdClient {
kv := pb.NewKVClient(conn)
lease := pb.NewLeaseClient(conn)
watch := pb.NewWatchClient(conn)
return &EtcdClient{conn, kv, lease, watch}
}
// RetryLeaseClient implements a LeaseClient that uses the client's FailFast retry policy.
func RetryLeaseClient(c *Client) pb.LeaseClient {
return &retryLeaseClient{pb.NewLeaseClient(c.conn), c.retryWrapper}
}