add put state batch

Signed-off-by: Fedor Partanskiy <[email protected]>

Author: pfi79

go.mod

@@ -20,3 +20,5 @@ require (
 	gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )
+
+replace github.com/hyperledger/fabric-protos-go-apiv2 => github.com/scientificideas/fabric-protos-go-apiv2 v0.0.0-20240819205904-84edd4adb1f8

go.sum

@@ -3,8 +3,6 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
 github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
-github.com/hyperledger/fabric-protos-go-apiv2 v0.3.3 h1:Xpd6fzG/KjAOHJsq7EQXY2l+qi/y8muxBaY7R6QWABk=
-github.com/hyperledger/fabric-protos-go-apiv2 v0.3.3/go.mod h1:2pq0ui6ZWA0cC8J+eCErgnMDCS1kPOEYVY+06ZAK0qE=
 github.com/kr/pretty v0.2.1 h1:Fmg33tUaq4/8ym9TJN1x7sLJnHVwhP33CNkpYV/7rwI=
 github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI=
 github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
@@ -13,6 +11,8 @@ github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
 github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/scientificideas/fabric-protos-go-apiv2 v0.0.0-20240819205904-84edd4adb1f8 h1:7hFAzj8zZyoNGo1zWIB5rYpY1GFkxwxKblcFsW8zQWY=
+github.com/scientificideas/fabric-protos-go-apiv2 v0.0.0-20240819205904-84edd4adb1f8/go.mod h1:1x9TFdg5b2M9vema7s4EdQZ1qUV4peAtF2jQj8FjLjc=
 github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
 github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
 golang.org/x/net v0.25.0 h1:d/OCCoBEUq33pjydKrGQhw7IlUPI2Oylr+8qLx49kac=

shim/handler.go

@@ -18,6 +18,10 @@ const (
 	created     state = "created"     // start state
 	established state = "established" // connection established
 	ready       state = "ready"       // ready for requests
+
+	defaultMaxSizePutStateBatch = 100
+	prefixMetaData              = "m"
+	prefixStateData             = "s"
 )
 
 // PeerChaincodeStream is the common stream interface for Peer - chaincode communication.
@@ -46,6 +50,11 @@ type Handler struct {
 	cc Chaincode
 	// state holds the current state of this handler.
 	state state
+	// if you can send the changes in batches.
+	usePutStateBatch     bool
+	maxSizePutStateBatch uint32
+	stateDataMutex       sync.RWMutex
+	stateData            map[string]map[string]*peer.StateKV
 
 	// Multiple queries (and one transaction) with different txids can be executing in parallel for this chaincode
 	// responseChannels is the channel on which responses are communicated by the shim to the chaincodeStub.
@@ -150,6 +159,7 @@ func newChaincodeHandler(peerChatStream PeerChaincodeStream, chaincode Chaincode
 		cc:               chaincode,
 		responseChannels: map[string]chan *peer.ChaincodeMessage{},
 		state:            created,
+		stateData:        map[string]map[string]*peer.StateKV{},
 	}
 }
 
@@ -188,6 +198,11 @@ func (h *Handler) handleInit(msg *peer.ChaincodeMessage) (*peer.ChaincodeMessage
 		return nil, fmt.Errorf("failed to marshal response: %s", err)
 	}
 
+	err = h.sendBatch(msg.ChannelId, msg.Txid)
+	if err != nil {
+		return nil, fmt.Errorf("failed send batch: %s", err)
+	}
+
 	return &peer.ChaincodeMessage{Type: peer.ChaincodeMessage_COMPLETED, Payload: resBytes, Txid: msg.Txid, ChaincodeEvent: stub.chaincodeEvent, ChannelId: stub.ChannelID}, nil
 }
 
@@ -214,6 +229,11 @@ func (h *Handler) handleTransaction(msg *peer.ChaincodeMessage) (*peer.Chaincode
 		return nil, fmt.Errorf("failed to marshal response: %s", err)
 	}
 
+	err = h.sendBatch(msg.ChannelId, msg.Txid)
+	if err != nil {
+		return nil, fmt.Errorf("failed send batch: %s", err)
+	}
+
 	return &peer.ChaincodeMessage{Type: peer.ChaincodeMessage_COMPLETED, Payload: resBytes, Txid: msg.Txid, ChaincodeEvent: stub.chaincodeEvent, ChannelId: stub.ChannelID}, nil
 }
 
@@ -312,6 +332,17 @@ func (h *Handler) handleGetStateMetadata(collection string, key string, channelI
 
 // handlePutState communicates with the peer to put state information into the ledger.
 func (h *Handler) handlePutState(collection string, key string, value []byte, channelID string, txid string) error {
+	if h.usePutStateBatch {
+		st := h.stateDataByID(channelID, txid)
+		st[prefixStateData+collection+key] = &peer.StateKV{
+			Key:        key,
+			Value:      value,
+			Collection: collection,
+			Type:       peer.StateKV_PUT_STATE,
+		}
+		return nil
+	}
+
 	// Construct payload for PUT_STATE
 	payloadBytes := marshalOrPanic(&peer.PutState{Collection: collection, Key: key, Value: value})
 
@@ -340,6 +371,19 @@ func (h *Handler) handlePutState(collection string, key string, value []byte, ch
 func (h *Handler) handlePutStateMetadataEntry(collection string, key string, metakey string, metadata []byte, channelID string, txID string) error {
 	// Construct payload for PUT_STATE_METADATA
 	md := &peer.StateMetadata{Metakey: metakey, Value: metadata}
+
+	if h.usePutStateBatch {
+		st := h.stateDataByID(channelID, txID)
+		st[prefixMetaData+collection+key] = &peer.StateKV{
+			Key:        key,
+			Collection: collection,
+			Metadata:   md,
+			Type:       peer.StateKV_PUT_STATE_METADATA,
+		}
+
+		return nil
+	}
+
 	payloadBytes := marshalOrPanic(&peer.PutStateMetadata{Collection: collection, Key: key, Metadata: md})
 
 	msg := &peer.ChaincodeMessage{Type: peer.ChaincodeMessage_PUT_STATE_METADATA, Payload: payloadBytes, Txid: txID, ChannelId: channelID}
@@ -365,6 +409,16 @@ func (h *Handler) handlePutStateMetadataEntry(collection string, key string, met
 
 // handleDelState communicates with the peer to delete a key from the state in the ledger.
 func (h *Handler) handleDelState(collection string, key string, channelID string, txid string) error {
+	if h.usePutStateBatch {
+		st := h.stateDataByID(channelID, txid)
+		st[prefixStateData+collection+key] = &peer.StateKV{
+			Key:        key,
+			Collection: collection,
+			Type:       peer.StateKV_DEL_STATE,
+		}
+		return nil
+	}
+
 	payloadBytes := marshalOrPanic(&peer.DelState{Collection: collection, Key: key})
 	msg := &peer.ChaincodeMessage{Type: peer.ChaincodeMessage_DEL_STATE, Payload: payloadBytes, Txid: txid, ChannelId: channelID}
 	// Execute the request and get response
@@ -388,6 +442,16 @@ func (h *Handler) handleDelState(collection string, key string, channelID string
 
 // handlePurgeState communicates with the peer to purge a state from private data
 func (h *Handler) handlePurgeState(collection string, key string, channelID string, txid string) error {
+	if h.usePutStateBatch {
+		st := h.stateDataByID(channelID, txid)
+		st[prefixStateData+collection+key] = &peer.StateKV{
+			Key:        key,
+			Collection: collection,
+			Type:       peer.StateKV_PURGE_PRIVATE_DATA,
+		}
+		return nil
+	}
+
 	payloadBytes := marshalOrPanic(&peer.DelState{Collection: collection, Key: key})
 	msg := &peer.ChaincodeMessage{Type: peer.ChaincodeMessage_PURGE_PRIVATE_DATA, Payload: payloadBytes, Txid: txid, ChannelId: channelID}
 	// Execute the request and get response
@@ -409,6 +473,69 @@ func (h *Handler) handlePurgeState(collection string, key string, channelID stri
 	return fmt.Errorf("[%s] incorrect chaincode message %s received. Expecting %s or %s", shorttxid(responseMsg.Txid), responseMsg.Type, peer.ChaincodeMessage_RESPONSE, peer.ChaincodeMessage_ERROR)
 }
 
+// handlePutStateBatch communicates with the peer to put state as batch all changes information into the ledger.
+func (h *Handler) handlePutStateBatch(batch *peer.PutStateBatch, channelID string, txid string) error {
+	// Construct payload for PUT_STATE_BATCH
+	payloadBytes := marshalOrPanic(batch)
+
+	msg := &peer.ChaincodeMessage{Type: peer.ChaincodeMessage_PUT_STATE_BATCH, Payload: payloadBytes, Txid: txid, ChannelId: channelID}
+
+	// Execute the request and get response
+	responseMsg, err := h.callPeerWithChaincodeMsg(msg, channelID, txid)
+	if err != nil {
+		return fmt.Errorf("[%s] error sending %s: %s", msg.Txid, peer.ChaincodeMessage_PUT_STATE_BATCH, err)
+	}
+
+	if responseMsg.Type == peer.ChaincodeMessage_RESPONSE {
+		// Success response
+		return nil
+	}
+
+	if responseMsg.Type == peer.ChaincodeMessage_ERROR {
+		// Error response
+		return fmt.Errorf("%s", responseMsg.Payload[:])
+	}
+
+	// Incorrect chaincode message received
+	return fmt.Errorf("[%s] incorrect chaincode message %s received. Expecting %s or %s", shorttxid(responseMsg.Txid), responseMsg.Type, peer.ChaincodeMessage_RESPONSE, peer.ChaincodeMessage_ERROR)
+}
+
+func (h *Handler) sendBatch(channelID string, txid string) error {
+	if !h.usePutStateBatch {
+		return nil
+	}
+
+	st := h.stateDataByID(channelID, txid)
+	txCtxID := transactionContextID(channelID, txid)
+
+	defer func() {
+		h.stateDataMutex.Lock()
+		delete(h.stateData, txCtxID)
+		h.stateDataMutex.Unlock()
+	}()
+
+	batch := &peer.PutStateBatch{}
+	for _, kv := range st {
+		batch.Kvs = append(batch.Kvs, kv)
+		if len(batch.Kvs) >= int(h.maxSizePutStateBatch) {
+			err := h.handlePutStateBatch(batch, channelID, txid)
+			if err != nil {
+				return fmt.Errorf("failed send state batch: %s", err)
+			}
+			batch.Kvs = batch.Kvs[:0]
+		}
+	}
+
+	if len(batch.Kvs) != 0 {
+		err := h.handlePutStateBatch(batch, channelID, txid)
+		if err != nil {
+			return fmt.Errorf("failed send state batch: %s", err)
+		}
+	}
+
+	return nil
+}
+
 func (h *Handler) handleGetStateByRange(collection, startKey, endKey string, metadata []byte,
 	channelID string, txid string) (*peer.QueryResponse, error) {
 	// Send GET_STATE_BY_RANGE message to peer chaincode support
@@ -655,6 +782,23 @@ func (h *Handler) handleEstablished(msg *peer.ChaincodeMessage) error {
 	}
 
 	h.state = ready
+	if len(msg.Payload) == 0 {
+		return nil
+	}
+
+	ccAdditionalParams := &peer.ChaincodeAdditionalParams{}
+	err := proto.Unmarshal(msg.Payload, ccAdditionalParams)
+	if err != nil {
+		return nil
+	}
+
+	h.usePutStateBatch = ccAdditionalParams.UsePutStateBatch
+	h.maxSizePutStateBatch = ccAdditionalParams.MaxSizePutStateBatch
+
+	if h.usePutStateBatch && h.maxSizePutStateBatch < defaultMaxSizePutStateBatch {
+		h.maxSizePutStateBatch = defaultMaxSizePutStateBatch
+	}
+
 	return nil
 }
 
@@ -697,6 +841,28 @@ func (h *Handler) handleMessage(msg *peer.ChaincodeMessage, errc chan error) err
 	return nil
 }
 
+func (h *Handler) stateDataByID(channelID string, txID string) map[string]*peer.StateKV {
+	txCtxID := transactionContextID(channelID, txID)
+
+	h.stateDataMutex.RLock()
+	st, ok := h.stateData[txCtxID]
+	h.stateDataMutex.RUnlock()
+	if ok {
+		return st
+	}
+
+	h.stateDataMutex.Lock()
+	defer h.stateDataMutex.Unlock()
+	st, ok = h.stateData[txCtxID]
+	if ok {
+		return st
+	}
+
+	st = make(map[string]*peer.StateKV)
+	h.stateData[txCtxID] = st
+	return st
+}
+
 // marshalOrPanic attempts to marshal the provided protobbuf message but will panic
 // when marshaling fails instead of returning an error.
 func marshalOrPanic(msg proto.Message) []byte {

shim/handler_test.go

@@ -49,6 +49,7 @@ func TestNewHandler_CreatedState(t *testing.T) {
 		cc:               cc,
 		responseChannels: map[string]chan *peer.ChaincodeMessage{},
 		state:            created,
+		stateData:        map[string]map[string]*peer.StateKV{},
 	}
 
 	handler := newChaincodeHandler(chatStream, cc)
@@ -125,7 +126,7 @@ func TestHandlerState(t *testing.T) {
 			}
 			err := handler.handleMessage(test.msg, nil)
 			if test.expectedErr != "" {
-				assert.Contains(t, err.Error(), test.expectedErr)
+				assert.ErrorContains(t, err, test.expectedErr)
 			} else {
 				assert.NoError(t, err)
 			}
@@ -217,7 +218,7 @@ func TestHandleMessage(t *testing.T) {
 
 			err := handler.handleMessage(test.msg, nil)
 			if test.expectedErr != "" {
-				assert.Contains(t, err.Error(), test.expectedErr)
+				assert.ErrorContains(t, err, test.expectedErr)
 			} else {
 				if err != nil {
 					t.Fatalf("Unexpected error for '%s': %s", test.name, err)
@@ -264,36 +265,36 @@ func TestHandlePeerCalls(t *testing.T) {
 	// force error by removing responseChannels
 	h.responseChannels = nil
 	_, err = h.handleGetState("col", "key", "channel", "txid")
-	assert.Contains(t, err.Error(), "[txid] error sending GET_STATE")
+	assert.ErrorContains(t, err, "[txid] error sending GET_STATE")
 
 	_, err = h.handleGetPrivateDataHash("col", "key", "channel", "txid")
-	assert.Contains(t, err.Error(), "[txid] error sending GET_PRIVATE_DATA_HASH")
+	assert.ErrorContains(t, err, "[txid] error sending GET_PRIVATE_DATA_HASH")
 
 	_, err = h.handleGetStateMetadata("col", "key", "channel", "txid")
-	assert.Contains(t, err.Error(), "[txid] error sending GET_STATE_METADATA")
+	assert.ErrorContains(t, err, "[txid] error sending GET_STATE_METADATA")
 
 	err = h.handlePutState("col", "key", []byte{}, "channel", "txid")
-	assert.Contains(t, err.Error(), "[txid] error sending PUT_STATE")
+	assert.ErrorContains(t, err, "[txid] error sending PUT_STATE")
 
 	err = h.handlePutStateMetadataEntry("col", "key", "mkey", []byte{}, "channel", "txid")
-	assert.Contains(t, err.Error(), "[txid] error sending PUT_STATE_METADATA")
+	assert.ErrorContains(t, err, "[txid] error sending PUT_STATE_METADATA")
 
 	err = h.handleDelState("col", "key", "channel", "txid")
-	assert.Contains(t, err.Error(), "[txid] error sending DEL_STATE")
+	assert.ErrorContains(t, err, "[txid] error sending DEL_STATE")
 
 	_, err = h.handleGetStateByRange("col", "start", "end", []byte{}, "channel", "txid")
-	assert.Contains(t, err.Error(), "[txid] error sending GET_STATE_BY_RANGE")
+	assert.ErrorContains(t, err, "[txid] error sending GET_STATE_BY_RANGE")
 
 	_, err = h.handleQueryStateNext("id", "channel", "txid")
-	assert.Contains(t, err.Error(), "cannot create response channel")
+	assert.ErrorContains(t, err, "cannot create response channel")
 
 	_, err = h.handleQueryStateClose("id", "channel", "txid")
-	assert.Contains(t, err.Error(), "cannot create response channel")
+	assert.ErrorContains(t, err, "cannot create response channel")
 
 	_, err = h.handleGetQueryResult("col", "query", []byte{}, "channel", "txid")
-	assert.Contains(t, err.Error(), "[txid] error sending GET_QUERY_RESULT")
+	assert.ErrorContains(t, err, "[txid] error sending GET_QUERY_RESULT")
 
 	_, err = h.handleGetHistoryForKey("key", "channel", "txid")
-	assert.Contains(t, err.Error(), "cannot create response channel")
+	assert.ErrorContains(t, err, "cannot create response channel")
 
 }