feat(requestmanager): read request from context (#381)

allow a request id to be set in the context for a request

Author: hannahhoward

graphsync.go

@@ -438,6 +438,10 @@ type RequestState uint64
 // RequestStates describe a set of request IDs and their current state
 type RequestStates map[RequestID]RequestState
 
+// RequestIDContextKey is used to the desired request id in context when
+// initializing a request
+type RequestIDContextKey struct{}
+
 const (
 	// Queued means a request has been received and is queued for processing
 	Queued RequestState = iota

requestmanager/client.go

@@ -178,9 +178,15 @@ func (rm *RequestManager) NewRequest(ctx context.Context,
 		return rm.singleErrorResponse(err)
 	}
 
+	requestID := graphsync.NewRequestID()
+	idFromContext := ctx.Value(graphsync.RequestIDContextKey{})
+	if existingRequestID, ok := idFromContext.(graphsync.RequestID); ok {
+		requestID = existingRequestID
+	}
+
 	inProgressRequestChan := make(chan inProgressRequest)
 
-	rm.send(&newRequestMessage{span, p, root, selectorNode, extensions, inProgressRequestChan}, ctx.Done())
+	rm.send(&newRequestMessage{requestID, span, p, root, selectorNode, extensions, inProgressRequestChan}, ctx.Done())
 	var receivedInProgressRequest inProgressRequest
 	select {
 	case <-rm.ctx.Done():

requestmanager/messages.go

@@ -104,6 +104,7 @@ func (trm *releaseRequestTaskMessage) handle(rm *RequestManager) {
 }
 
 type newRequestMessage struct {
+	requestID             graphsync.RequestID
 	span                  trace.Span
 	p                     peer.ID
 	root                  ipld.Link
@@ -115,7 +116,7 @@ type newRequestMessage struct {
 func (nrm *newRequestMessage) handle(rm *RequestManager) {
 	var ipr inProgressRequest
 
-	ipr.request, ipr.incoming, ipr.incomingError = rm.newRequest(nrm.span, nrm.p, nrm.root, nrm.selector, nrm.extensions)
+	ipr.request, ipr.incoming, ipr.incomingError = rm.newRequest(nrm.requestID, nrm.span, nrm.p, nrm.root, nrm.selector, nrm.extensions)
 	ipr.requestID = ipr.request.ID()
 
 	select {

requestmanager/requestmanager_test.go

@@ -988,6 +988,23 @@ func TestStats(t *testing.T) {
 	require.Len(t, peerState.Pending, 0)
 }
 
+func TestCaptureRequestIDFromContext(t *testing.T) {
+	ctx := context.Background()
+	td := newTestData(ctx, t)
+
+	requestCtx, cancel := context.WithTimeout(ctx, time.Second)
+	defer cancel()
+	expectedID := graphsync.NewRequestID()
+	requestCtx = context.WithValue(requestCtx, graphsync.RequestIDContextKey{}, expectedID)
+
+	peers := testutil.GeneratePeers(1)
+
+	_, _ = td.requestManager.NewRequest(requestCtx, peers[0], td.blockChain.TipLink, td.blockChain.Selector())
+
+	requestRecords := readNNetworkRequests(requestCtx, t, td, 1)
+	require.Equal(t, expectedID, requestRecords[0].gsr.ID())
+}
+
 type requestRecord struct {
 	gsr gsmsg.GraphSyncRequest
 	p   peer.ID

requestmanager/server.go

@@ -60,8 +60,7 @@ func (rm *RequestManager) cleanupInProcessRequests() {
 	}
 }
 
-func (rm *RequestManager) newRequest(parentSpan trace.Span, p peer.ID, root ipld.Link, selector ipld.Node, extensions []graphsync.ExtensionData) (gsmsg.GraphSyncRequest, chan graphsync.ResponseProgress, chan error) {
-	requestID := graphsync.NewRequestID()
+func (rm *RequestManager) newRequest(requestID graphsync.RequestID, parentSpan trace.Span, p peer.ID, root ipld.Link, selector ipld.Node, extensions []graphsync.ExtensionData) (gsmsg.GraphSyncRequest, chan graphsync.ResponseProgress, chan error) {
 
 	parentSpan.SetAttributes(attribute.String("requestID", requestID.String()))
 	ctx, span := otel.Tracer("graphsync").Start(trace.ContextWithSpan(rm.ctx, parentSpan), "newRequest")