From 5b52f8778920fafec1d808f63fa6bf094e2d8ff5 Mon Sep 17 00:00:00 2001
From: Aleksandr Bezobchuk <alexanderbez@users.noreply.github.com>
Date: Wed, 17 Feb 2021 12:03:11 -0500
Subject: [PATCH] ABCI: Fix ReCheckTx for Socket Client (#6124)

Fixes the race condition between a callback being set and called during ReCheckTx. Note, I do not see equivalent logic in the gRPC client (anymore) as #5439 suggests, so only the socket client was updated.

closes: #5439
---
 CHANGELOG_PENDING.md         |  1 +
 abci/client/client.go        | 60 +++++++++++++++++++++---------------
 abci/client/socket_client.go |  4 +--
 3 files changed, 38 insertions(+), 27 deletions(-)

diff --git a/CHANGELOG_PENDING.md b/CHANGELOG_PENDING.md
index db980b8f6..b85218d67 100644
--- a/CHANGELOG_PENDING.md
+++ b/CHANGELOG_PENDING.md
@@ -66,6 +66,7 @@ Friendly reminder, we have a [bug bounty program](https://hackerone.com/tendermi
 
 ### BUG FIXES
 
+- [ABCI] \#6124 Fixes a panic condition during callback execution in `ReCheckTx` during high tx load. (@alexanderbez)
 - [types] \#5523 Change json naming of `PartSetHeader` within `BlockID` from `parts` to `part_set_header` (@marbar3778)
 - [privval] \#5638 Increase read/write timeout to 5s and calculate ping interval based on it (@JoeKash)
 - [blockchain/v1] [\#5701](https://github.com/tendermint/tendermint/pull/5701) Handle peers without blocks (@melekes)
diff --git a/abci/client/client.go b/abci/client/client.go
index 3b55aebae..4cc5dafff 100644
--- a/abci/client/client.go
+++ b/abci/client/client.go
@@ -80,12 +80,8 @@ func NewClient(addr, transport string, mustConnect bool) (client Client, err err
 	return
 }
 
-//----------------------------------------
-
 type Callback func(*types.Request, *types.Response)
 
-//----------------------------------------
-
 type ReqRes struct {
 	*types.Request
 	*sync.WaitGroup
@@ -107,34 +103,50 @@ func NewReqRes(req *types.Request) *ReqRes {
 	}
 }
 
-// Sets the callback for this ReqRes atomically.
-// If reqRes is already done, calls cb immediately.
-// NOTE: reqRes.cb should not change if reqRes.done.
-// NOTE: only one callback is supported.
-func (reqRes *ReqRes) SetCallback(cb func(res *types.Response)) {
-	reqRes.mtx.Lock()
+// Sets sets the callback. If reqRes is already done, it will call the cb
+// immediately. Note, reqRes.cb should not change if reqRes.done and only one
+// callback is supported.
+func (r *ReqRes) SetCallback(cb func(res *types.Response)) {
+	r.mtx.Lock()
 
-	if reqRes.done {
-		reqRes.mtx.Unlock()
-		cb(reqRes.Response)
+	if r.done {
+		r.mtx.Unlock()
+		cb(r.Response)
 		return
 	}
 
-	reqRes.cb = cb
-	reqRes.mtx.Unlock()
+	r.cb = cb
+	r.mtx.Unlock()
 }
 
-func (reqRes *ReqRes) GetCallback() func(*types.Response) {
-	reqRes.mtx.Lock()
-	defer reqRes.mtx.Unlock()
-	return reqRes.cb
+// InvokeCallback invokes a thread-safe execution of the configured callback
+// if non-nil.
+func (r *ReqRes) InvokeCallback() {
+	r.mtx.Lock()
+	defer r.mtx.Unlock()
+
+	if r.cb != nil {
+		r.cb(r.Response)
+	}
 }
 
-// NOTE: it should be safe to read reqRes.cb without locks after this.
-func (reqRes *ReqRes) SetDone() {
-	reqRes.mtx.Lock()
-	reqRes.done = true
-	reqRes.mtx.Unlock()
+// GetCallback returns the configured callback of the ReqRes object which may be
+// nil. Note, it is not safe to concurrently call this in cases where it is
+// marked done and SetCallback is called before calling GetCallback as that
+// will invoke the callback twice and create a potential race condition.
+//
+// ref: https://github.com/tendermint/tendermint/issues/5439
+func (r *ReqRes) GetCallback() func(*types.Response) {
+	r.mtx.Lock()
+	defer r.mtx.Unlock()
+	return r.cb
+}
+
+// SetDone marks the ReqRes object as done.
+func (r *ReqRes) SetDone() {
+	r.mtx.Lock()
+	r.done = true
+	r.mtx.Unlock()
 }
 
 func waitGroup1() (wg *sync.WaitGroup) {
diff --git a/abci/client/socket_client.go b/abci/client/socket_client.go
index 5c3fec937..0569f4c6b 100644
--- a/abci/client/socket_client.go
+++ b/abci/client/socket_client.go
@@ -226,9 +226,7 @@ func (cli *socketClient) didRecvResponse(res *types.Response) error {
 	//
 	// NOTE: It is possible this callback isn't set on the reqres object. At this
 	// point, in which case it will be called after, when it is set.
-	if cb := reqres.GetCallback(); cb != nil {
-		cb(res)
-	}
+	reqres.InvokeCallback()
 
 	return nil
 }