From 9d94d7d4789583a09b9e0765290cca9f9c670bb3 Mon Sep 17 00:00:00 2001
From: siddontang <siddontang@gmail.com>
Date: Thu, 6 Sep 2018 17:21:59 +0800
Subject: [PATCH] use same process ID expect meeting unknown response (#18)

---
 db/tidb/bank.go           |  9 ++++++++-
 pkg/control/control.go    | 17 +++++++++++++----
 pkg/core/client.go        |  8 +++++---
 pkg/model/cas_register.go |  8 ++++++++
 pkg/model/register.go     |  8 ++++++++
 5 files changed, 42 insertions(+), 8 deletions(-)

diff --git a/db/tidb/bank.go b/db/tidb/bank.go
index da69b0a..b103e67 100644
--- a/db/tidb/bank.go
+++ b/db/tidb/bank.go
@@ -191,6 +191,13 @@ type bankResponse struct {
 	Unknown bool
 }
 
+var _ core.UnknownResponse = (*bankResponse)(nil)
+
+// IsUnknown implements UnknownResponse interface
+func (r bankResponse) IsUnknown() bool {
+	return r.Unknown
+}
+
 func newBankEvent(v interface{}, id uint) porcupine.Event {
 	if _, ok := v.(bankRequest); ok {
 		return porcupine.Event{Kind: porcupine.CallEvent, Value: v, Id: id}
@@ -408,7 +415,7 @@ func mergeTransferEvents(index int, events tsoEvents, e *tsoEvent) (tsoEvents, e
 	curBalance, _ := e.GetBalances(index)
 
 	if !checkBalance(index, events, curBalance) {
-		return nil, fmt.Errorf("invalid event %s", e)
+		return nil, fmt.Errorf("%d %v invalid event %s", index, events, e)
 	}
 
 	events = append(events, e)
diff --git a/pkg/control/control.go b/pkg/control/control.go
index 6e8cd7f..dad76ec 100644
--- a/pkg/control/control.go
+++ b/pkg/control/control.go
@@ -33,7 +33,8 @@ type Controller struct {
 	ctx    context.Context
 	cancel context.CancelFunc
 
-	proc int64
+	proc         int64
+	requestCount int64
 
 	recorder *history.Recorder
 }
@@ -180,9 +181,8 @@ func (c *Controller) onClientLoop(i int) {
 	ctx, cancel := context.WithTimeout(c.ctx, c.cfg.RunTime)
 	defer cancel()
 
-	for i := 0; i < c.cfg.RequestCount; i++ {
-		procID := atomic.AddInt64(&c.proc, 1)
-
+	procID := atomic.AddInt64(&c.proc, 1)
+	for atomic.AddInt64(&c.requestCount, 1) <= int64(c.cfg.RequestCount) {
 		request := client.NextRequest()
 
 		if err := c.recorder.RecordRequest(procID, request); err != nil {
@@ -190,11 +190,20 @@ func (c *Controller) onClientLoop(i int) {
 		}
 
 		response := client.Invoke(ctx, node, request)
+		isUnknown := true
+		if v, ok := response.(core.UnknownResponse); ok {
+			isUnknown = v.IsUnknown()
+		}
 
 		if err := c.recorder.RecordResponse(procID, response); err != nil {
 			log.Fatalf("record response %v failed %v", response, err)
 		}
 
+		// If Unknown, we need to use another process ID.
+		if isUnknown {
+			procID = atomic.AddInt64(&c.proc, 1)
+		}
+
 		select {
 		case <-ctx.Done():
 			return
diff --git a/pkg/core/client.go b/pkg/core/client.go
index 98956ba..0b2351f 100644
--- a/pkg/core/client.go
+++ b/pkg/core/client.go
@@ -4,9 +4,10 @@ import (
 	"context"
 )
 
-// RequestGenerator generates a request.
-type RequestGenerator interface {
-	Generate() interface{}
+// UnknownResponse means we don't know wether this operation
+// succeeds or not.
+type UnknownResponse interface {
+	IsUnknown() bool
 }
 
 // Client applies the request to the database.
@@ -18,6 +19,7 @@ type Client interface {
 	// TearDown tears down the client.
 	TearDown(ctx context.Context, nodes []string, node string) error
 	// Invoke invokes a request to the database.
+	// Mostly, the return Response should implement UnknownResponse interface
 	Invoke(ctx context.Context, node string, r interface{}) interface{}
 	// NextRequest generates a request for latter Invoke.
 	NextRequest() interface{}
diff --git a/pkg/model/cas_register.go b/pkg/model/cas_register.go
index b2b3a23..2057616 100644
--- a/pkg/model/cas_register.go
+++ b/pkg/model/cas_register.go
@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 
 	"github.com/anishathalye/porcupine"
+	"github.com/siddontang/chaos/pkg/core"
 	"github.com/siddontang/chaos/pkg/history"
 )
 
@@ -32,6 +33,13 @@ type CasRegisterResponse struct {
 	Unknown bool // used when operation times out
 }
 
+var _ core.UnknownResponse = (*CasRegisterResponse)(nil)
+
+// IsUnknown implements UnknownResponse interface
+func (r CasRegisterResponse) IsUnknown() bool {
+	return r.Unknown
+}
+
 // CasRegisterModel returns a cas register model
 func CasRegisterModel() porcupine.Model {
 	return porcupine.Model{
diff --git a/pkg/model/register.go b/pkg/model/register.go
index eb3a34f..e9ec066 100644
--- a/pkg/model/register.go
+++ b/pkg/model/register.go
@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 
 	"github.com/anishathalye/porcupine"
+	"github.com/siddontang/chaos/pkg/core"
 	"github.com/siddontang/chaos/pkg/history"
 )
 
@@ -29,6 +30,13 @@ type RegisterResponse struct {
 	Value   int
 }
 
+var _ core.UnknownResponse = (*RegisterResponse)(nil)
+
+// IsUnknown implements UnknownResponse interface
+func (r RegisterResponse) IsUnknown() bool {
+	return r.Unknown
+}
+
 // RegisterModel returns a read/write register model
 func RegisterModel() porcupine.Model {
 	return porcupine.Model{