Fix the client connection is closing issue in robustness test

Signed-off-by: Chun-Hung Tseng <henrytseng@google.com>

Author: henrybear327

tests/robustness/client/client.go

@@ -30,6 +30,10 @@ import (
 	"go.etcd.io/etcd/tests/v3/robustness/report"
 )
 
+var (
+	errClientIsClosed = errors.New("client is closed")
+)
+
 // RecordingClient provides a semi-etcd client (different interface than
 // clientv3.Client) that records all the requests and responses made. Doesn't
 // allow for concurrent requests to conform to model.AppendableHistory requirements.
@@ -45,6 +49,9 @@ type RecordingClient struct {
 	// mux ensures order of request appending.
 	kvMux        sync.Mutex
 	kvOperations *model.AppendableHistory
+
+	isClosed bool
+	sync.Mutex
 }
 
 var _ clientv3.KV = (*RecordingClient)(nil)
@@ -69,11 +76,28 @@ func NewRecordingClient(endpoints []string, ids identity.Provider, baseTime time
 		client:       *cc,
 		kvOperations: model.NewAppendableHistory(ids),
 		baseTime:     baseTime,
+		isClosed:     false,
 	}, nil
 }
 
 func (c *RecordingClient) Close() error {
-	return c.client.Close()
+	c.Lock()
+	defer c.Unlock()
+	if c.IsClosed() {
+		return nil
+	}
+
+	err := c.client.Close()
+	if err != nil {
+		c.isClosed = true
+	}
+	return err
+}
+
+func (c *RecordingClient) IsClosed() bool {
+	c.Lock()
+	defer c.Unlock()
+	return c.isClosed
 }
 
 func (c *RecordingClient) Report() report.ClientReport {
@@ -89,11 +113,23 @@ func (c *RecordingClient) Do(ctx context.Context, op clientv3.Op) (clientv3.OpRe
 }
 
 func (c *RecordingClient) Get(ctx context.Context, key string, opts ...clientv3.OpOption) (*clientv3.GetResponse, error) {
+	c.Lock()
+	defer c.Unlock()
+	if c.isClosed {
+		return nil, errClientIsClosed
+	}
+
 	op := clientv3.OpGet(key, opts...)
 	return c.Range(ctx, key, string(op.RangeBytes()), op.Rev(), op.Limit())
 }
 
 func (c *RecordingClient) Range(ctx context.Context, start, end string, revision, limit int64) (*clientv3.GetResponse, error) {
+	c.Lock()
+	defer c.Unlock()
+	if c.isClosed {
+		return nil, errClientIsClosed
+	}
+
 	ops := []clientv3.OpOption{}
 	if end != "" {
 		ops = append(ops, clientv3.WithRange(end))
@@ -114,6 +150,12 @@ func (c *RecordingClient) Range(ctx context.Context, start, end string, revision
 }
 
 func (c *RecordingClient) Put(ctx context.Context, key, value string, _ ...clientv3.OpOption) (*clientv3.PutResponse, error) {
+	c.Lock()
+	defer c.Unlock()
+	if c.isClosed {
+		return nil, errClientIsClosed
+	}
+
 	c.kvMux.Lock()
 	defer c.kvMux.Unlock()
 	callTime := time.Since(c.baseTime)
@@ -124,6 +166,12 @@ func (c *RecordingClient) Put(ctx context.Context, key, value string, _ ...clien
 }
 
 func (c *RecordingClient) Delete(ctx context.Context, key string, _ ...clientv3.OpOption) (*clientv3.DeleteResponse, error) {
+	c.Lock()
+	defer c.Unlock()
+	if c.isClosed {
+		return nil, errClientIsClosed
+	}
+
 	c.kvMux.Lock()
 	defer c.kvMux.Unlock()
 	callTime := time.Since(c.baseTime)
@@ -172,10 +220,22 @@ func (w *wrappedTxn) Commit() (*clientv3.TxnResponse, error) {
 }
 
 func (c *RecordingClient) Txn(ctx context.Context) clientv3.Txn {
+	c.Lock()
+	defer c.Unlock()
+	if c.isClosed {
+		return nil
+	}
+
 	return &wrappedTxn{txn: c.client.Txn(ctx), c: c}
 }
 
 func (c *RecordingClient) LeaseGrant(ctx context.Context, ttl int64) (*clientv3.LeaseGrantResponse, error) {
+	c.Lock()
+	defer c.Unlock()
+	if c.isClosed {
+		return nil, errClientIsClosed
+	}
+
 	c.kvMux.Lock()
 	defer c.kvMux.Unlock()
 	callTime := time.Since(c.baseTime)
@@ -186,6 +246,12 @@ func (c *RecordingClient) LeaseGrant(ctx context.Context, ttl int64) (*clientv3.
 }
 
 func (c *RecordingClient) LeaseRevoke(ctx context.Context, leaseID int64) (*clientv3.LeaseRevokeResponse, error) {
+	c.Lock()
+	defer c.Unlock()
+	if c.isClosed {
+		return nil, errClientIsClosed
+	}
+
 	c.kvMux.Lock()
 	defer c.kvMux.Unlock()
 	callTime := time.Since(c.baseTime)
@@ -196,6 +262,12 @@ func (c *RecordingClient) LeaseRevoke(ctx context.Context, leaseID int64) (*clie
 }
 
 func (c *RecordingClient) PutWithLease(ctx context.Context, key string, value string, leaseID int64) (*clientv3.PutResponse, error) {
+	c.Lock()
+	defer c.Unlock()
+	if c.isClosed {
+		return nil, errClientIsClosed
+	}
+
 	opts := clientv3.WithLease(clientv3.LeaseID(leaseID))
 	c.kvMux.Lock()
 	defer c.kvMux.Unlock()
@@ -207,6 +279,12 @@ func (c *RecordingClient) PutWithLease(ctx context.Context, key string, value st
 }
 
 func (c *RecordingClient) Defragment(ctx context.Context) (*clientv3.DefragmentResponse, error) {
+	c.Lock()
+	defer c.Unlock()
+	if c.isClosed {
+		return nil, errClientIsClosed
+	}
+
 	c.kvMux.Lock()
 	defer c.kvMux.Unlock()
 	callTime := time.Since(c.baseTime)
@@ -217,6 +295,12 @@ func (c *RecordingClient) Defragment(ctx context.Context) (*clientv3.DefragmentR
 }
 
 func (c *RecordingClient) Compact(ctx context.Context, rev int64, _ ...clientv3.CompactOption) (*clientv3.CompactResponse, error) {
+	c.Lock()
+	defer c.Unlock()
+	if c.isClosed {
+		return nil, errClientIsClosed
+	}
+
 	c.kvMux.Lock()
 	defer c.kvMux.Unlock()
 	callTime := time.Since(c.baseTime)
@@ -227,48 +311,90 @@ func (c *RecordingClient) Compact(ctx context.Context, rev int64, _ ...clientv3.
 }
 
 func (c *RecordingClient) MemberList(ctx context.Context, opts ...clientv3.OpOption) (*clientv3.MemberListResponse, error) {
+	c.Lock()
+	defer c.Unlock()
+	if c.isClosed {
+		return nil, errClientIsClosed
+	}
+
 	c.kvMux.Lock()
 	defer c.kvMux.Unlock()
 	resp, err := c.client.MemberList(ctx, opts...)
 	return resp, err
 }
 
 func (c *RecordingClient) MemberAdd(ctx context.Context, peerAddrs []string) (*clientv3.MemberAddResponse, error) {
+	c.Lock()
+	defer c.Unlock()
+	if c.isClosed {
+		return nil, errClientIsClosed
+	}
+
 	c.kvMux.Lock()
 	defer c.kvMux.Unlock()
 	resp, err := c.client.MemberAdd(ctx, peerAddrs)
 	return resp, err
 }
 
 func (c *RecordingClient) MemberAddAsLearner(ctx context.Context, peerAddrs []string) (*clientv3.MemberAddResponse, error) {
+	c.Lock()
+	defer c.Unlock()
+	if c.isClosed {
+		return nil, errClientIsClosed
+	}
+
 	c.kvMux.Lock()
 	defer c.kvMux.Unlock()
 	resp, err := c.client.MemberAddAsLearner(ctx, peerAddrs)
 	return resp, err
 }
 
 func (c *RecordingClient) MemberRemove(ctx context.Context, id uint64) (*clientv3.MemberRemoveResponse, error) {
+	c.Lock()
+	defer c.Unlock()
+	if c.isClosed {
+		return nil, errClientIsClosed
+	}
+
 	c.kvMux.Lock()
 	defer c.kvMux.Unlock()
 	resp, err := c.client.MemberRemove(ctx, id)
 	return resp, err
 }
 
 func (c *RecordingClient) MemberUpdate(ctx context.Context, id uint64, peerAddrs []string) (*clientv3.MemberUpdateResponse, error) {
+	c.Lock()
+	defer c.Unlock()
+	if c.isClosed {
+		return nil, errClientIsClosed
+	}
+
 	c.kvMux.Lock()
 	defer c.kvMux.Unlock()
 	resp, err := c.client.MemberUpdate(ctx, id, peerAddrs)
 	return resp, err
 }
 
 func (c *RecordingClient) MemberPromote(ctx context.Context, id uint64) (*clientv3.MemberPromoteResponse, error) {
+	c.Lock()
+	defer c.Unlock()
+	if c.isClosed {
+		return nil, errClientIsClosed
+	}
+
 	c.kvMux.Lock()
 	defer c.kvMux.Unlock()
 	resp, err := c.client.MemberPromote(ctx, id)
 	return resp, err
 }
 
 func (c *RecordingClient) Status(ctx context.Context, endpoint string) (*clientv3.StatusResponse, error) {
+	c.Lock()
+	defer c.Unlock()
+	if c.isClosed {
+		return nil, errClientIsClosed
+	}
+
 	c.kvMux.Lock()
 	defer c.kvMux.Unlock()
 	resp, err := c.client.Status(ctx, endpoint)
@@ -280,6 +406,12 @@ func (c *RecordingClient) Endpoints() []string {
 }
 
 func (c *RecordingClient) Watch(ctx context.Context, key string, rev int64, withPrefix bool, withProgressNotify bool, withPrevKV bool) clientv3.WatchChan {
+	c.Lock()
+	defer c.Unlock()
+	if c.isClosed {
+		return nil
+	}
+
 	request := model.WatchRequest{
 		Key:                key,
 		Revision:           rev,
@@ -333,6 +465,12 @@ func (c *RecordingClient) watch(ctx context.Context, request model.WatchRequest)
 }
 
 func (c *RecordingClient) RequestProgress(ctx context.Context) error {
+	c.Lock()
+	defer c.Unlock()
+	if c.isClosed {
+		return errClientIsClosed
+	}
+
 	return c.client.RequestProgress(ctx)
 }
 
@@ -434,7 +572,7 @@ func (cs *ClientSet) close() {
 		return
 	}
 	for _, c := range cs.clients {
-		c.Close()
+		_ = c.Close()
 	}
 	cs.closed = true
 }

tests/robustness/client/watch.go

@@ -103,6 +103,9 @@ resetWatch:
 			return nil
 		}
 		watch := c.Watch(ctx, "", lastRevision+1, true, true, false)
+		if watch == nil {
+			return nil
+		}
 		for {
 			select {
 			case revision, ok := <-maxRevisionChan:
@@ -161,13 +164,19 @@ func openWatchPeriodically(ctx context.Context, g *errgroup.Group, c *RecordingC
 		g.Go(func() error {
 			resp, err := c.Get(ctx, "/key")
 			if err != nil {
+				if errors.Is(err, errClientIsClosed) {
+					return nil
+				}
 				return err
 			}
 			rev := resp.Header.Revision + backgroundWatchConfig.RevisionOffset
 
 			watchCtx, cancel := context.WithCancel(ctx)
 			defer cancel()
 			w := c.Watch(watchCtx, "", rev, true, true, true)
+			if w == nil {
+				return nil
+			}
 			for {
 				select {
 				case <-ctx.Done():