metering: refactor wasm input bytes metering to a function in the met…

…ering package
streamingfast · Oct 14, 2024 · ced4b37 · ced4b37
1 parent fc843ca
commit ced4b37
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 1 deletion.
diff --git a/metering/metering.go b/metering/metering.go
@@ -60,6 +60,10 @@ func WithBytesMeteringOptions(meter dmetering.Meter, logger *zap.Logger) []dstor
 	return opts
 }
 
+func AddWasmInputBytes(ctx context.Context, n int) {
+	dmetering.GetBytesMeter(ctx).CountInc(MeterWasmInputBytes, n)
+}
+
 func GetTotalBytesRead(meter dmetering.Meter) uint64 {
 	total := uint64(meter.GetCount(TotalReadBytes))
 	return total

diff --git a/metering/metering_test.go b/metering/metering_test.go
@@ -293,6 +293,15 @@ func TestLiveSourceMiddlewareHandlerFactory(t *testing.T) {
 	}
 }
 
+func TestAddWasmInputBytes(t *testing.T) {
+	ctx := dmetering.WithBytesMeter(context.Background())
+	meter := dmetering.GetBytesMeter(ctx)
+
+	AddWasmInputBytes(ctx, 10)
+
+	assert.Equal(t, 10, meter.GetCount(MeterWasmInputBytes))
+}
+
 func TestSend(t *testing.T) {
 	ctx := dmetering.WithBytesMeter(context.Background())
 	meter := dmetering.GetBytesMeter(ctx)

diff --git a/pipeline/process_block.go b/pipeline/process_block.go
@@ -256,7 +256,8 @@ func (p *Pipeline) handleStepNew(ctx context.Context, clock *pbsubstreams.Clock,
 		return fmt.Errorf("pre block hook: %w", err)
 	}
 
-	dmetering.GetBytesMeter(ctx).CountInc(metering.MeterWasmInputBytes, execOutput.Len())
+	metering.AddWasmInputBytes(ctx, execOutput.Len())
+
 	if err := p.executeModules(ctx, execOutput); err != nil {
 		return fmt.Errorf("execute modules: %w", err)
 	}