From 3606a6236ee791f447712242448418f814ab4757 Mon Sep 17 00:00:00 2001
From: Angelo De Caro <adc@zurich.ibm.com>
Date: Thu, 12 Dec 2024 07:24:10 +0100
Subject: [PATCH] use struct's context

Signed-off-by: Angelo De Caro <adc@zurich.ibm.com>
---
 platform/fabric/chaincode.go                        | 5 +++++
 platform/fabric/core/generic/chaincode/discovery.go | 8 +++++++-
 platform/fabric/core/generic/chaincode/invoke.go    | 8 ++++++--
 platform/fabric/driver/chaincode.go                 | 1 +
 4 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/platform/fabric/chaincode.go b/platform/fabric/chaincode.go
index 46a2e9a1d..c75640c08 100644
--- a/platform/fabric/chaincode.go
+++ b/platform/fabric/chaincode.go
@@ -146,6 +146,11 @@ func (i *ChaincodeDiscover) WithFilterByMSPIDs(mspIDs ...string) *ChaincodeDisco
 	return i
 }
 
+func (i *ChaincodeDiscover) WithContext(context context.Context) *ChaincodeDiscover {
+	i.ChaincodeDiscover.WithContext(context)
+	return i
+}
+
 type ChaincodeInvocation struct {
 	driver.ChaincodeInvocation
 }
diff --git a/platform/fabric/core/generic/chaincode/discovery.go b/platform/fabric/core/generic/chaincode/discovery.go
index ef04a65d7..e187de33d 100644
--- a/platform/fabric/core/generic/chaincode/discovery.go
+++ b/platform/fabric/core/generic/chaincode/discovery.go
@@ -29,6 +29,7 @@ type Discovery struct {
 	QueryForPeers       bool
 
 	DefaultTTL time.Duration
+	Context    context.Context
 }
 
 func NewDiscovery(chaincode *Chaincode) *Discovery {
@@ -36,6 +37,7 @@ func NewDiscovery(chaincode *Chaincode) *Discovery {
 	return &Discovery{
 		chaincode:  chaincode,
 		DefaultTTL: chaincode.ChannelConfig.DiscoveryDefaultTTLS(),
+		Context:    context.Background(),
 	}
 }
 
@@ -231,7 +233,7 @@ func (d *Discovery) query(req *discovery.Request) (discovery.Response, error) {
 		ClientIdentity:    signerRaw,
 		ClientTlsCertHash: ClientTLSCertHash,
 	}
-	timeout, cancel := context.WithTimeout(context.Background(), d.chaincode.ChannelConfig.DiscoveryTimeout())
+	timeout, cancel := context.WithTimeout(d.Context, d.chaincode.ChannelConfig.DiscoveryTimeout())
 	defer cancel()
 	cl, err := pc.DiscoveryClient()
 	if err != nil {
@@ -319,6 +321,10 @@ func (d *Discovery) ChaincodeVersion() (string, error) {
 	return "", errors.Errorf("chaincode [%s] not found", d.chaincode.name)
 }
 
+func (d *Discovery) WithContext(context context.Context) {
+	d.Context = context
+}
+
 func ccCall(ccNames ...string) []*peer.ChaincodeCall {
 	var call []*peer.ChaincodeCall
 	for _, ccName := range ccNames {
diff --git a/platform/fabric/core/generic/chaincode/invoke.go b/platform/fabric/core/generic/chaincode/invoke.go
index 51c283892..8863ce70c 100644
--- a/platform/fabric/core/generic/chaincode/invoke.go
+++ b/platform/fabric/core/generic/chaincode/invoke.go
@@ -63,6 +63,7 @@ func NewInvoke(chaincode *Chaincode, newChaincodeDiscover NewChaincodeDiscoverFu
 		NumRetries:           int(chaincode.NumRetries),
 		RetrySleep:           chaincode.RetrySleep,
 		NewChaincodeDiscover: newChaincodeDiscover,
+		Context:              context.Background(),
 	}
 }
 
@@ -316,6 +317,8 @@ func (i *Invoke) prepare(query bool) (string, *pb.Proposal, []*pb.ProposalRespon
 			i.EndorsersMSPIDs...,
 		).WithImplicitCollections(
 			i.ImplicitCollectionMSPIDs...,
+		).WithContext(
+			i.Context,
 		)
 		if query {
 			discovery.WithForQuery()
@@ -473,7 +476,8 @@ func (i *Invoke) collectResponses(endorserClients []pb.EndorserClient, signedPro
 	for _, endorser := range endorserClients {
 		go func(endorser pb.EndorserClient) {
 			defer wg.Done()
-			proposalResp, err := endorser.ProcessProposal(context.Background(), signedProposal)
+			// TODO: we could evaluate the policy already here after we get a result to see if still need more answers
+			proposalResp, err := endorser.ProcessProposal(i.Context, signedProposal)
 			if err != nil {
 				errorCh <- err
 				return
@@ -553,7 +557,7 @@ func (i *Invoke) broadcast(txID string, env *common.Envelope) error {
 	if err := i.Chaincode.Broadcaster.Broadcast(i.Context, env); err != nil {
 		return err
 	}
-	return i.Chaincode.Finality.IsFinal(context.Background(), txID)
+	return i.Chaincode.Finality.IsFinal(i.Context, txID)
 }
 
 func (i *Invoke) checkQueryPolicy(errs []error, successes int, n int) error {
diff --git a/platform/fabric/driver/chaincode.go b/platform/fabric/driver/chaincode.go
index 2868e933d..58cc20e84 100644
--- a/platform/fabric/driver/chaincode.go
+++ b/platform/fabric/driver/chaincode.go
@@ -92,6 +92,7 @@ type ChaincodeDiscover interface {
 	WithImplicitCollections(mspIDs ...string) ChaincodeDiscover
 	WithForQuery() ChaincodeDiscover
 	ChaincodeVersion() (string, error)
+	WithContext(ctx context.Context)
 }
 
 // Chaincode exposes chaincode-related functions