From 68fd863c8778a19ce7538cc2b7410c74d7c42f4c Mon Sep 17 00:00:00 2001
From: Marek Siarkowicz <siarkowicz@google.com>
Date: Thu, 1 Dec 2022 20:30:31 +0100
Subject: [PATCH] tests: Mark failed requests as timed out

Signed-off-by: Marek Siarkowicz <siarkowicz@google.com>
---
 tests/linearizability/client.go               | 53 ++++++++++++++++++-
 tests/linearizability/linearizability_test.go |  2 +-
 2 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/tests/linearizability/client.go b/tests/linearizability/client.go
index b8cf2dd03..2bfb6a079 100644
--- a/tests/linearizability/client.go
+++ b/tests/linearizability/client.go
@@ -25,9 +25,13 @@ import (
 
 type recordingClient struct {
 	client clientv3.Client
-	id     int
+
+	// id of the next write operation. If needed a new id might be requested from idProvider.
+	id         int
+	idProvider idProvider
 
 	operations []porcupine.Operation
+	failed     []porcupine.Operation
 }
 
 func NewClient(endpoints []string, ids idProvider) (*recordingClient, error) {
@@ -43,7 +47,9 @@ func NewClient(endpoints []string, ids idProvider) (*recordingClient, error) {
 	return &recordingClient{
 		client:     *cc,
 		id:         ids.ClientId(),
+		idProvider: ids,
 		operations: []porcupine.Operation{},
+		failed:     []porcupine.Operation{},
 	}, nil
 }
 
@@ -76,6 +82,19 @@ func (c *recordingClient) Put(ctx context.Context, key, value string) error {
 	callTime := time.Now()
 	resp, err := c.client.Put(ctx, key, value)
 	returnTime := time.Now()
+	if err != nil {
+		c.failed = append(c.failed, porcupine.Operation{
+			ClientId: c.id,
+			Input:    etcdRequest{op: Put, key: key, putData: value},
+			Call:     callTime.UnixNano(),
+			Output:   etcdResponse{err: err},
+			Return:   0, // For failed writes we don't know when request has really finished.
+		})
+		// Operations of single client needs to be sequential.
+		// As we don't know return time of failed operations, all new writes need to be done with new client id.
+		c.id = c.idProvider.ClientId()
+		return err
+	}
 	var revision int64
 	if resp != nil && resp.Header != nil {
 		revision = resp.Header.Revision
@@ -94,6 +113,19 @@ func (c *recordingClient) Delete(ctx context.Context, key string) error {
 	callTime := time.Now()
 	resp, err := c.client.Delete(ctx, key)
 	returnTime := time.Now()
+	if err != nil {
+		c.failed = append(c.failed, porcupine.Operation{
+			ClientId: c.id,
+			Input:    etcdRequest{op: Delete, key: key},
+			Call:     callTime.UnixNano(),
+			Output:   etcdResponse{err: err},
+			Return:   0, // For failed writes we don't know when request has really finished.
+		})
+		// Operations of single client needs to be sequential.
+		// As we don't know return time of failed operations, all new writes need to be done with new client id.
+		c.id = c.idProvider.ClientId()
+		return err
+	}
 	var revision int64
 	var deleted int64
 	if resp != nil && resp.Header != nil {
@@ -109,3 +141,22 @@ func (c *recordingClient) Delete(ctx context.Context, key string) error {
 	})
 	return nil
 }
+
+func (c *recordingClient) Operations() []porcupine.Operation {
+	operations := make([]porcupine.Operation, 0, len(c.operations)+len(c.failed))
+	var maxTime int64
+	for _, op := range c.operations {
+		operations = append(operations, op)
+		if op.Return > maxTime {
+			maxTime = op.Return
+		}
+	}
+	for _, op := range c.failed {
+		if op.Call > maxTime {
+			continue
+		}
+		op.Return = maxTime + 1
+		operations = append(operations, op)
+	}
+	return operations
+}
diff --git a/tests/linearizability/linearizability_test.go b/tests/linearizability/linearizability_test.go
index ae5539e9d..570cb881c 100644
--- a/tests/linearizability/linearizability_test.go
+++ b/tests/linearizability/linearizability_test.go
@@ -162,7 +162,7 @@ func simulateTraffic(ctx context.Context, t *testing.T, clus *e2e.EtcdProcessClu
 
 			config.traffic.Run(ctx, c, limiter, ids)
 			mux.Lock()
-			operations = append(operations, c.operations...)
+			operations = append(operations, c.Operations()...)
 			mux.Unlock()
 		}(c)
 	}