metering: MetricsSender.Send reads identity from the context

metering/metering.go

@@ -8,6 +8,7 @@ import (
 
 	"github.com/streamingfast/bstream"
 	pbbstream "github.com/streamingfast/bstream/pb/sf/bstream/v1"
+	"github.com/streamingfast/dauth"
 	"github.com/streamingfast/dmetering"
 	"github.com/streamingfast/dstore"
 	"github.com/streamingfast/substreams/metrics"
@@ -149,7 +150,8 @@ func NewMetricsSender() *MetricsSender {
 	}
 }
 
-func (ms *MetricsSender) Send(ctx context.Context, organizationID, apiKeyID, ip, userMeta, outputModuleHash, endpoint string) {
+// Send reads the identity from ctx on every call, so headers refreshed mid-stream apply.
+func (ms *MetricsSender) Send(ctx context.Context, outputModuleHash, endpoint string) {
 	ms.Lock()
 	defer ms.Unlock()
 
@@ -200,11 +202,12 @@ func (ms *MetricsSender) Send(ctx context.Context, organizationID, apiKeyID, ip,
 		return true
 	})
 
+	auth := dauth.FromContext(ctx)
 	event := dmetering.Event{
-		OrganizationID:   organizationID,
-		ApiKeyID:         apiKeyID,
-		IpAddress:        ip,
-		Meta:             userMeta,
+		OrganizationID:   auth.OrganizationID(),
+		ApiKeyID:         auth.APIKeyID(),
+		IpAddress:        auth.RealIP(),
+		Meta:             auth.Meta(),
 		OutputModuleHash: outputModuleHash,
 
 		Endpoint:  endpoint,

metering/metering_test.go

@@ -12,6 +12,7 @@ import (
 
 	"github.com/streamingfast/bstream"
 	pbbstream "github.com/streamingfast/bstream/pb/sf/bstream/v1"
+	"github.com/streamingfast/dauth"
 	"github.com/streamingfast/dmetering"
 	"github.com/streamingfast/dstore"
 	pbsubstreamstest "github.com/streamingfast/substreams/pb/sf/substreams/v1/test"
@@ -390,8 +391,8 @@ func TestSend_externalCalls(t *testing.T) {
 	AddExternalCalls(ctx, "eth_getBalance", 4)
 
 	sender := NewMetricsSender()
-	sender.Send(ctx, "org1", "apiKey1", "127.0.0.1", "meta", "outputModuleHash", "endpoint")
-	sender.Send(ctx, "org1", "apiKey1", "127.0.0.1", "meta", "outputModuleHash", "endpoint")
+	sender.Send(ctx, "outputModuleHash", "endpoint")
+	sender.Send(ctx, "outputModuleHash", "endpoint")
 
 	assert.Len(t, emitter.events, 2)
 	assert.Equal(t, float64(5), emitter.events[0].Metrics["external_calls_eth_call"])
@@ -428,7 +429,7 @@ func TestSend(t *testing.T) {
 	outputModuleHash := "outputModuleHash"
 
 	AddEgressBytes(ctx, proto.Size(resp))
-	metericsSender.Send(ctx, "org1", "apiKey1", "127.0.0.1", "meta", outputModuleHash, "endpoint")
+	metericsSender.Send(withTestIdentity(ctx), outputModuleHash, "endpoint")
 
 	// Verify the emitted event
 	assert.Len(t, emitter.events, 1)
@@ -491,7 +492,7 @@ func TestSendParallel(t *testing.T) {
 
 			time.Sleep(time.Duration(randomInt()) * time.Nanosecond)
 			AddEgressBytes(ctx, proto.Size(resp))
-			metricsSender.Send(ctx, "org1", "apiKey1", "127.0.0.1", "meta", "outputModuleHash", "endpoint")
+			metricsSender.Send(ctx, "outputModuleHash", "endpoint")
 		}()
 	}
 
@@ -567,3 +568,24 @@ func (m *mockEmitter) Emit(ctx context.Context, event dmetering.Event) {
 func (m *mockEmitter) Shutdown(err error) {
 	return
 }
+
+func withTestIdentity(ctx context.Context) context.Context {
+	return dauth.WithTrustedHeaders(ctx, dauth.TrustedHeaders{
+		dauth.HeaderOrganizationID: "org1",
+		dauth.HeaderApiKeyID:       "apiKey1",
+		dauth.HeaderIP:             "127.0.0.1",
+		dauth.HeaderMeta:           "meta",
+	})
+}
+
+func TestSend_noTrustedHeaders(t *testing.T) {
+	ctx := dmetering.WithBytesMeter(context.Background())
+	emitter := &mockEmitter{}
+	ctx = reqctx.WithEmitter(ctx, emitter)
+
+	NewMetricsSender().Send(ctx, "outputModuleHash", "endpoint")
+
+	require.Len(t, emitter.events, 1)
+	assert.Equal(t, "", emitter.events[0].OrganizationID)
+	assert.Equal(t, "", emitter.events[0].Meta)
+}

pipeline/terminate.go

@@ -8,7 +8,6 @@ import (
 	"time"
 
 	"github.com/streamingfast/bstream/stream"
-	"github.com/streamingfast/dauth"
 	"go.uber.org/zap"
 
 	"github.com/streamingfast/substreams/metering"
@@ -62,14 +61,9 @@ func (p *Pipeline) OnStreamTerminated(ctx context.Context, err error) error {
 
 	if reqctx.Details(ctx).IsTier2Request {
 		// metrics on tier2 are just sent at the end of the ProcessBlock
-		auth := dauth.FromContext(ctx)
-		organizationID := auth.OrganizationID()
-		apiKeyID := auth.APIKeyID()
-		userMeta := auth.Meta()
-		ip := auth.RealIP()
 		outputModuleHash := reqctx.OutputModuleHash(ctx)
 		metricsSender := metering.GetMetricsSender(ctx)
-		metricsSender.Send(ctx, organizationID, apiKeyID, ip, userMeta, outputModuleHash, "sf.substreams.internal.v2/ProcessRange")
+		metricsSender.Send(ctx, outputModuleHash, "sf.substreams.internal.v2/ProcessRange")
 		err := p.returnInternalModuleProgressOutputs(p.lastFinalClock, true)
 		if err != nil {
 			logger.Error("returning internal module progress outputs", zap.Error(err))

service/tier1.go

@@ -1284,12 +1284,6 @@ func tier1ResponseHandler(
 	debugOutputForModules []string,
 	supportBuffering bool,
 ) substreams.ResponseFunc {
-	auth := dauth.FromContext(ctx)
-	organizationID := auth.OrganizationID()
-	apiKeyID := auth.APIKeyID()
-	userMeta := auth.Meta()
-	ip := auth.RealIP()
-
 	outputModuleHash := reqctx.OutputModuleHash(ctx)
 
 	endpoint := "sf.substreams.rpc.v2/Blocks"
@@ -1350,7 +1344,8 @@ func tier1ResponseHandler(
 		egressBytes := proto.Size(respAny.(proto.Message))
 		begin := time.Now()
 		if err := streamSrv.SendMsg(respAny); err != nil {
-			logger.Info("unable to send block probably due to client disconnecting", zap.String("user_id", organizationID), zap.String("api_key_id", apiKeyID), zap.Error(err))
+			auth := dauth.FromContext(ctx)
+			logger.Info("unable to send block probably due to client disconnecting", zap.String("user_id", auth.OrganizationID()), zap.String("api_key_id", auth.APIKeyID()), zap.Error(err))
 			return connect.NewError(connect.CodeUnavailable, err)
 		}
 		stats.RecordReadTime(begin)
@@ -1366,7 +1361,7 @@ func tier1ResponseHandler(
 		stats.RecordEgress(egressBytes)
 		metering.AddEgressBytes(ctx, egressBytes)
 
-		metericsSender.Send(ctx, organizationID, apiKeyID, ip, userMeta, outputModuleHash, endpoint)
+		metericsSender.Send(ctx, outputModuleHash, endpoint)
 		return nil
 	}
 }

test/collector_test.go

@@ -71,7 +71,7 @@ func (c *responseCollector) Collect(respAny substreams.ResponseFromAnyTier) erro
 	case *pbsubstreamsrpc.Response:
 		c.responses = append(c.responses, resp)
 		metering.AddEgressBytes(c.ctx, proto.Size(resp))
-		c.sender.Send(c.ctx, "test_org", "test_api_key", "10.0.0.1", "test_meta", "testOutputHash", "tier1")
+		c.sender.Send(c.ctx, "testOutputHash", "tier1")
 	case *pbssinternal.ProcessRangeResponse:
 		// in non-test code, this is 'passed through' from tier2 to tier1 to the user as a pbsubstreamsrpc.response
 		if blockScopedData := resp.GetBlockScopedData(); blockScopedData != nil {
@@ -91,7 +91,7 @@ func (c *responseCollector) Collect(respAny substreams.ResponseFromAnyTier) erro
 		}
 
 		c.internalResponses = append(c.internalResponses, resp)
-		c.sender.Send(c.ctx, "test_org", "test_api_key", "10.0.0.1", "test_meta", "testOutputHash", "tier2")
+		c.sender.Send(c.ctx, "testOutputHash", "tier2")
 	}
 	return nil
 }