add output module hash to metering

streamingfast · Oct 30, 2024 · ebdc718 · ebdc718
1 parent c9d10c6
commit ebdc718
Show file tree

Hide file tree

Showing 8 changed files with 44 additions and 19 deletions.
diff --git a/go.mod b/go.mod
@@ -57,7 +57,7 @@ require (
 	github.com/rs/cors v1.10.0
 	github.com/schollz/closestmatch v2.1.0+incompatible
 	github.com/shopspring/decimal v1.3.1
-	github.com/streamingfast/dmetering v0.0.0-20240816165719-51768d3da951
+	github.com/streamingfast/dmetering v0.0.0-20241028183059-d11d4ec85e05
 	github.com/streamingfast/dmetrics v0.0.0-20230919161904-206fa8ebd545
 	github.com/streamingfast/sf-tracing v0.0.0-20240430173521-888827872b90
 	github.com/streamingfast/shutter v1.5.0

diff --git a/go.sum b/go.sum
@@ -533,8 +533,8 @@ github.com/streamingfast/derr v0.0.0-20230515163924-8570aaa43fe1 h1:xJB7rXnOHLes
 github.com/streamingfast/derr v0.0.0-20230515163924-8570aaa43fe1/go.mod h1:QSm/AfaDsE0k1xBYi0lW580YJ/WDV/FKZI628tkZR0Y=
 github.com/streamingfast/dgrpc v0.0.0-20240219152146-57bb131c39ca h1:/k5H6MUo5Vi8AKPsSr+TMeA/XJ0uMyEX6feHpOozTlQ=
 github.com/streamingfast/dgrpc v0.0.0-20240219152146-57bb131c39ca/go.mod h1:NuKCwOHjbT0nRji0O+7+c70AiBfLHEKNoovs/gFfMPY=
-github.com/streamingfast/dmetering v0.0.0-20240816165719-51768d3da951 h1:6o6MS3JHrp9A7V6EBHbR7W7mzVCFmXc8U0AjTfvz7PI=
-github.com/streamingfast/dmetering v0.0.0-20240816165719-51768d3da951/go.mod h1:UqWuX3REU/IInBUaymFN2eLjuvz+/0SsoUFjeQlLNyI=
+github.com/streamingfast/dmetering v0.0.0-20241028183059-d11d4ec85e05 h1:jjx6kO2z0mNSKElL8YHnjt65+NB/CSsJH5C9efyrzw8=
+github.com/streamingfast/dmetering v0.0.0-20241028183059-d11d4ec85e05/go.mod h1:UqWuX3REU/IInBUaymFN2eLjuvz+/0SsoUFjeQlLNyI=
 github.com/streamingfast/dmetrics v0.0.0-20230919161904-206fa8ebd545 h1:SUl04bZKGAv207lp7/6CHOJIRpjUKunwItrno3K463Y=
 github.com/streamingfast/dmetrics v0.0.0-20230919161904-206fa8ebd545/go.mod h1:JbxEDbzWRG1dHdNIPrYfuPllEkktZMgm40AwVIBENcw=
 github.com/streamingfast/dstore v0.1.1-0.20241011152904-9acd6205dc14 h1:/2HxIOzAgUBKyxjDO4IJPzBBaEAtzwipb/2/JGsOArA=

diff --git a/metering/metering.go b/metering/metering.go
@@ -121,6 +121,8 @@ func (ms *MetricsSender) Send(ctx context.Context, userID, apiKeyID, ip, userMet
 		endpoint = fmt.Sprintf("%s%s", endpoint, "Backfill")
 	}
 
+	outputModuleHash := reqctx.OutputModuleHash(ctx)
+
 	meter := dmetering.GetBytesMeter(ctx)
 
 	bytesRead := meter.BytesReadDelta()
@@ -143,10 +145,11 @@ func (ms *MetricsSender) Send(ctx context.Context, userID, apiKeyID, ip, userMet
 	meter.CountInc(TotalWriteBytes, int(totalWriteBytes))
 
 	event := dmetering.Event{
-		UserID:    userID,
-		ApiKeyID:  apiKeyID,
-		IpAddress: ip,
-		Meta:      userMeta,
+		UserID:           userID,
+		ApiKeyID:         apiKeyID,
+		IpAddress:        ip,
+		Meta:             userMeta,
+		OutputModuleHash: outputModuleHash,
 
 		Endpoint: endpoint,
 		Metrics: map[string]float64{

diff --git a/reqctx/context.go b/reqctx/context.go
@@ -23,6 +23,8 @@ var tracerKey = contextKeyType(2)
 var spanKey = contextKeyType(3)
 var reqStatsKey = contextKeyType(4)
 var moduleExecutionTracingConfigKey = contextKeyType(5)
+var outputModuleHashKey = contextKeyType(6)
+var tier2RequestParametersKeyKey = contextKeyType(7)
 
 func Logger(ctx context.Context) *zap.Logger {
 	return logging.Logger(ctx, zap.NewNop())

diff --git a/reqctx/metering.go b/reqctx/metering.go
@@ -31,3 +31,17 @@ func IsBackfillerRequest(ctx context.Context) bool {
 	_, ok = md[backFillerKey]
 	return ok
 }
+
+type outputModuleKeyType int
+
+func WithOutputModuleHash(ctx context.Context, hash string) context.Context {
+	return context.WithValue(ctx, outputModuleHashKey, hash)
+}
+
+func OutputModuleHash(ctx context.Context) string {
+	hash, ok := ctx.Value(outputModuleHashKey).(string)
+	if !ok {
+		return ""
+	}
+	return hash
+}
diff --git a/reqctx/tier2request.go b/reqctx/tier2request.go
@@ -16,10 +16,6 @@ type Tier2RequestParameters struct {
 	WASMModules map[string]string
 }
 
-type tier2RequestParametersKey int
-
-const tier2RequestParametersKeyKey = tier2RequestParametersKey(0)
-
 func WithTier2RequestParameters(ctx context.Context, parameters Tier2RequestParameters) context.Context {
 	return context.WithValue(ctx, tier2RequestParametersKeyKey, parameters)
 }

diff --git a/service/tier1.go b/service/tier1.go
@@ -219,24 +219,25 @@ func (s *Tier1Service) Blocks(
 		mut.Unlock()
 	}()
 
-	respFunc := tier1ResponseHandler(respContext, &mut, logger, stream)
-
-	span.SetAttributes(attribute.Int64("substreams.tier", 1))
-
 	request := req.Msg
 	if request.Modules == nil {
 		return connect.NewError(connect.CodeInvalidArgument, fmt.Errorf("missing modules in request"))
 	}
 
-	if err := ValidateTier1Request(request, s.blockType); err != nil {
-		return connect.NewError(connect.CodeInvalidArgument, fmt.Errorf("validate request: %w", err))
-	}
-
 	execGraph, err := exec.NewOutputModuleGraph(request.OutputModule, request.ProductionMode, request.Modules, bstream.GetProtocolFirstStreamableBlock)
 	if err != nil {
 		return bsstream.NewErrInvalidArg(err.Error())
 	}
 	outputModuleHash := execGraph.ModuleHashes().Get(request.OutputModule)
+	ctx = reqctx.WithOutputModuleHash(ctx, outputModuleHash)
+
+	respFunc := tier1ResponseHandler(respContext, &mut, logger, stream)
+
+	span.SetAttributes(attribute.Int64("substreams.tier", 1))
+
+	if err := ValidateTier1Request(request, s.blockType); err != nil {
+		return connect.NewError(connect.CodeInvalidArgument, fmt.Errorf("validate request: %w", err))
+	}
 
 	moduleNames := make([]string, len(request.Modules.Modules))
 	for i := 0; i < len(moduleNames); i++ {

diff --git a/service/tier2.go b/service/tier2.go
@@ -9,6 +9,8 @@ import (
 	"sync"
 	"time"
 
+	"github.com/streamingfast/bstream"
+
 	"connectrpc.com/connect"
 	"github.com/RoaringBitmap/roaring/roaring64"
 	"github.com/streamingfast/bstream/stream"
@@ -206,6 +208,13 @@ func (s *Tier2Service) ProcessRange(request *pbssinternal.ProcessRangeRequest, s
 		return connect.NewError(connect.CodeInvalidArgument, fmt.Errorf("validate request: %w", err))
 	}
 
+	execGraph, err := exec.NewOutputModuleGraph(request.OutputModule, true, request.Modules, bstream.GetProtocolFirstStreamableBlock)
+	if err != nil {
+		return bsstream.NewErrInvalidArg(err.Error())
+	}
+	outputModuleHash := execGraph.ModuleHashes().Get(request.OutputModule)
+	ctx = reqctx.WithOutputModuleHash(ctx, outputModuleHash)
+
 	emitter, err := dmetering.New(request.MeteringConfig, logger)
 	if err != nil {
 		return connect.NewError(connect.CodeInternal, fmt.Errorf("unable to initialize dmetering: %w", err))