Running req extractor (#75)

* Add data layer interfaces.

Signed-off-by: Mohammad <mohammad.nassar@ibm.com>

* add NotificationSource in datalayer for asynchronous event pipeline.

Signed-off-by: Mohammad <mohammad.nassar@ibm.com>

* Remove unneeded starttime.

Signed-off-by: Mohammad <mohammad.nassar@ibm.com>

* Add running request counter extractor.

Signed-off-by: Mohammad <mohammad.nassar@ibm.com>

* Align with model interface changes.

Signed-off-by: Mohammad <mohammad.nassar@ibm.com>

* Add plugin factory.

Signed-off-by: Mohammad <mohammad.nassar@ibm.com>

* Add factory to source plugin.

Signed-off-by: Mohammad <mohammad.nassar@ibm.com>

* Use datastore without handle.

Signed-off-by: Mohammad <mohammad.nassar@ibm.com>

* Restore handle file.

Signed-off-by: Mohammad <mohammad.nassar@ibm.com>

* Simplify notification source.

Signed-off-by: Mohammad <mohammad.nassar@ibm.com>

* Fix doc comments.

Signed-off-by: Mohammad <mohammad.nassar@ibm.com>

* Simplify the extractor.

Signed-off-by: Mohammad <mohammad.nassar@ibm.com>

* Register the plugin.

Signed-off-by: Mohammad <mohammad.nassar@ibm.com>

* Remove lock.

Signed-off-by: Mohammad <mohammad.nassar@ibm.com>

* Adding the factory and register the extractor in runner.

Signed-off-by: Mohammad <mohammad.nassar@ibm.com>

* Revert notification source changes.

Signed-off-by: Mohammad <mohammad.nassar@ibm.com>

* Small fix.

Signed-off-by: Mohammad <mohammad.nassar@ibm.com>

* Fix comments.

Signed-off-by: Mohammad Nassar <mohammad.nassar@ibm.com>

* Small fix.

Signed-off-by: Mohammad <mohammad.nassar@ibm.com>

* Add unittest file.

Signed-off-by: Mohammad <mohammad.nassar@ibm.com>

* Fix lint error.

Signed-off-by: Mohammad <mohammad.nassar@ibm.com>

* Update after datalayer change.

Signed-off-by: Mohammad <mohammad.nassar@ibm.com>

* Fix datastore lint error.

Signed-off-by: Mohammad <mohammad.nassar@ibm.com>

* Change factory name.

Signed-off-by: Mohammad <mohammad.nassar@ibm.com>

* Revert datastore changes.

Signed-off-by: Mohammad <mohammad.nassar@ibm.com>

* Fix conflicts.

Signed-off-by: Mohammad <mohammad.nassar@ibm.com>

* Fix conflicts.

Signed-off-by: Mohammad <mohammad.nassar@ibm.com>

* Fix conflicts.

Signed-off-by: Mohammad <mohammad.nassar@ibm.com>

---------

Signed-off-by: Mohammad <mohammad.nassar@ibm.com>
Signed-off-by: Mohammad-nassar10 <79787844+Mohammad-nassar10@users.noreply.github.com>
Signed-off-by: Mohammad Nassar <mohammad.nassar@ibm.com>

Author: Mohammad-nassar10

cmd/runner/runner.go

@@ -38,10 +38,12 @@ import (
 	logutil "github.com/llm-d/llm-d-inference-payload-processor/pkg/common/observability/logging"
 	"github.com/llm-d/llm-d-inference-payload-processor/pkg/common/observability/profiling"
 	"github.com/llm-d/llm-d-inference-payload-processor/pkg/common/observability/tracing"
+	"github.com/llm-d/llm-d-inference-payload-processor/pkg/datastore"
 	"github.com/llm-d/llm-d-inference-payload-processor/pkg/framework"
 	"github.com/llm-d/llm-d-inference-payload-processor/pkg/metrics"
 	"github.com/llm-d/llm-d-inference-payload-processor/pkg/plugins/basemodelextractor"
 	"github.com/llm-d/llm-d-inference-payload-processor/pkg/plugins/bodyfieldtoheader"
+	inflightrequests "github.com/llm-d/llm-d-inference-payload-processor/pkg/plugins/datalayer/inflightrequests"
 	notificationsource "github.com/llm-d/llm-d-inference-payload-processor/pkg/plugins/datalayer/notificationsource"
 	runserver "github.com/llm-d/llm-d-inference-payload-processor/pkg/server"
 	"github.com/llm-d/llm-d-inference-payload-processor/version"
@@ -172,6 +174,8 @@ func (r *Runner) Run(ctx context.Context) error {
 
 	handle := framework.NewHandle(ctx, mgr)
 
+	ds := datastore.NewStore()
+
 	// Register factories for all known in-tree plugins
 	r.registerInTreePlugins()
 
@@ -220,6 +224,18 @@ func (r *Runner) Run(ctx context.Context) error {
 		}
 	}
 
+	// Wire the inflight-requests data pipeline: extractor → notification source.
+	// TODO: config-driven path does not yet support NotificationSource + extractors.
+	notifSrc, err := notificationsource.New("default", inflightrequests.NewInflightRequestsExtractor(ds))
+	if err != nil {
+		setupLog.Error(err, "failed to create notification source")
+		return err
+	}
+	if err := notifSrc.Start(ctx); err != nil {
+		setupLog.Error(err, "failed to start notification source")
+		return err
+	}
+
 	// Setup ExtProc Server Runner.
 	serverRunner := &runserver.ExtProcServerRunner{
 		GrpcPort:        opts.GRPCPort,
@@ -253,6 +269,7 @@ func (r *Runner) Run(ctx context.Context) error {
 func (r *Runner) registerInTreePlugins() {
 	framework.Register(bodyfieldtoheader.BodyFieldToHeaderPluginType, bodyfieldtoheader.BodyFieldToHeaderPluginFactory)
 	framework.Register(basemodelextractor.BaseModelToHeaderPluginType, basemodelextractor.BaseModelToHeaderPluginFactory)
+	framework.Register(inflightrequests.PluginType, inflightrequests.ExtractorFactory)
 	framework.Register(notificationsource.PluginType, notificationsource.Factory)
 }
 

pkg/plugins/datalayer/inflightrequests/plugin.go

@@ -0,0 +1,136 @@
+/*
+Copyright 2026 The llm-d Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package inflightrequests
+
+import (
+	"context"
+	"encoding/json"
+
+	"github.com/llm-d/llm-d-inference-payload-processor/pkg/framework"
+	"github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/datalayer"
+)
+
+const (
+	// PluginType is the identifier used when registering this extractor.
+	PluginType = "inflight-requests-extractor"
+
+	// InflightRequestsAttributeKey is the attribute key written to each model's attribute store.
+	InflightRequestsAttributeKey = "inflight-requests"
+)
+
+// compile-time interface assertion
+var _ datalayer.Extractor = &InflightRequestsExtractor{}
+
+// ExtractorFactory creates a InflightRequestsExtractor with a nil DataStore.
+// The factory path is limited: the DataStore is not available via framework.Handle,
+// so the created extractor cannot write to the store. Use NewInflightRequestsExtractor
+// directly when constructing for production use.
+func ExtractorFactory(name string, _ json.RawMessage, _ framework.Handle) (framework.Plugin, error) {
+	return NewInflightRequestsExtractor(nil).WithName(name), nil
+}
+
+// InflightRequestsCount holds in-flight request and token counts for one model.
+type InflightRequestsCount struct {
+	Requests int64
+	Tokens   int64
+}
+
+func (r InflightRequestsCount) Clone() datalayer.Cloneable { return r }
+
+// InflightRequestsExtractor tracks in-flight request counts and token sums per model.
+// It writes InflightRequestsCount to each model's InflightRequestsAttributeKey attribute.
+//
+// Extract is assumed to be called from a single goroutine (the NotificationSource event loop).
+// If parallel dispatch is introduced, add a sync.Mutex around counters and the DataStore write.
+//
+// TODO: counters leak if a request fails without a corresponding ResponseEventType (e.g. connection
+// drop, upstream error, context cancellation). The call site should fire a
+// synthetic ResponseEventType in its error/EOF path to keep counts accurate.
+type InflightRequestsExtractor struct {
+	typedName framework.TypedName
+	dataStore datalayer.DataStore
+	counters  map[string]InflightRequestsCount
+}
+
+func NewInflightRequestsExtractor(ds datalayer.DataStore) *InflightRequestsExtractor {
+	return &InflightRequestsExtractor{
+		typedName: framework.TypedName{Type: PluginType, Name: PluginType},
+		dataStore: ds,
+		counters:  make(map[string]InflightRequestsCount),
+	}
+}
+
+func (e *InflightRequestsExtractor) TypedName() framework.TypedName { return e.typedName }
+
+// WithName sets the instance name, used by the factory when the plugin is configured by name.
+func (e *InflightRequestsExtractor) WithName(name string) *InflightRequestsExtractor {
+	e.typedName.Name = name
+	return e
+}
+
+func (e *InflightRequestsExtractor) Extract(_ context.Context, events []datalayer.Event) error {
+	updated := map[string]InflightRequestsCount{}
+
+	for _, ev := range events {
+		switch ev.Type {
+		case datalayer.RequestEventType:
+			p, ok := ev.Payload.(datalayer.RequestPayload)
+			if !ok {
+				continue
+			}
+			model, _ := p.Request.Body["model"].(string)
+			if model == "" {
+				continue
+			}
+			maxTokens, _ := p.Request.Body["max_tokens"].(float64)
+			c := e.counters[model]
+			c.Requests++
+			c.Tokens += int64(maxTokens)
+			e.counters[model] = c
+			updated[model] = c
+
+		case datalayer.ResponseEventType:
+			p, ok := ev.Payload.(datalayer.ResponsePayload)
+			if !ok {
+				continue
+			}
+			model, _ := p.Request.Body["model"].(string)
+			if model == "" {
+				continue
+			}
+			maxTokens, _ := p.Request.Body["max_tokens"].(float64)
+			c := e.counters[model]
+			floorDecrement(&c.Requests, 1)
+			floorDecrement(&c.Tokens, int64(maxTokens))
+			e.counters[model] = c
+			updated[model] = c
+		}
+	}
+
+	for model, c := range updated {
+		e.dataStore.GetOrCreateModel(model).GetAttributes().Put(InflightRequestsAttributeKey, c)
+	}
+	return nil
+}
+
+// floorDecrement decrements v by delta, flooring at zero.
+func floorDecrement(v *int64, delta int64) {
+	*v -= delta
+	if *v < 0 {
+		*v = 0
+	}
+}

pkg/plugins/datalayer/inflightrequests/plugin_test.go

@@ -0,0 +1,210 @@
+/*
+Copyright 2026 The llm-d Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package inflightrequests
+
+import (
+	"context"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/llm-d/llm-d-inference-payload-processor/pkg/framework"
+	"github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/datalayer"
+)
+
+// fakeDataStore is an in-memory DataStore for tests.
+type fakeDataStore struct {
+	mu     sync.Mutex
+	models map[string]datalayer.Model
+}
+
+func newFakeDataStore() *fakeDataStore {
+	return &fakeDataStore{models: make(map[string]datalayer.Model)}
+}
+
+func (f *fakeDataStore) GetOrCreateModel(name string) datalayer.Model {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	if m, ok := f.models[name]; ok {
+		return m
+	}
+	m := datalayer.NewModel(name)
+	f.models[name] = m
+	return m
+}
+
+// makeRequestEvent creates a RequestEventType event with model and max_tokens.
+func makeRequestEvent(model string, maxTokens float64) datalayer.Event {
+	req := framework.NewInferenceRequest()
+	req.Body["model"] = model
+	req.Body["max_tokens"] = maxTokens
+	return datalayer.Event{
+		Type:    datalayer.RequestEventType,
+		Payload: datalayer.RequestPayload{Request: req},
+	}
+}
+
+// makeResponseEvent creates a ResponseEventType event with model, duration, and max_tokens.
+// maxTokens mirrors the original request's max_tokens so the extractor can decrement correctly.
+func makeResponseEvent(model string, durationMs int, maxTokens float64) datalayer.Event {
+	req := framework.NewInferenceRequest()
+	req.Body["model"] = model
+	req.Body["max_tokens"] = maxTokens
+	return datalayer.Event{
+		Type: datalayer.ResponseEventType,
+		Payload: datalayer.ResponsePayload{
+			Request:  req,
+			Response: framework.NewInferenceResponse(),
+			Duration: time.Duration(durationMs) * time.Millisecond,
+		},
+	}
+}
+
+// getInflightRequests asserts the inflight-requests attribute exists for model and returns it.
+func getInflightRequests(t testing.TB, ds *fakeDataStore, model string) InflightRequestsCount {
+	t.Helper()
+	val, ok := ds.GetOrCreateModel(model).GetAttributes().Get(InflightRequestsAttributeKey)
+	if !ok {
+		t.Fatalf("expected %q attribute for model %q", InflightRequestsAttributeKey, model)
+	}
+	rc, ok := val.(InflightRequestsCount)
+	if !ok {
+		t.Fatalf("expected InflightRequestsCount for model %q", model)
+	}
+	return rc
+}
+
+func newInflightRequestsTest(t *testing.T) (*InflightRequestsExtractor, *fakeDataStore) {
+	t.Helper()
+	ds := newFakeDataStore()
+	return NewInflightRequestsExtractor(ds), ds
+}
+
+func TestRequestIncrementsCounter(t *testing.T) {
+	ext, ds := newInflightRequestsTest(t)
+
+	batch := []datalayer.Event{makeRequestEvent("m1", 100)}
+	if err := ext.Extract(context.Background(), batch); err != nil {
+		t.Fatalf("Extract failed: %v", err)
+	}
+
+	rc := getInflightRequests(t, ds, "m1")
+	if rc.Requests != 1 {
+		t.Errorf("expected Requests=1, got %d", rc.Requests)
+	}
+	if rc.Tokens != 100 {
+		t.Errorf("expected Tokens=100, got %d", rc.Tokens)
+	}
+}
+
+func TestResponseDecrementsCounter(t *testing.T) {
+	ext, ds := newInflightRequestsTest(t)
+
+	// Response carries the original request's max_tokens so the extractor can decrement correctly.
+	batch := []datalayer.Event{
+		makeRequestEvent("m1", 100),
+		makeResponseEvent("m1", 50, 100),
+	}
+	if err := ext.Extract(context.Background(), batch); err != nil {
+		t.Fatalf("Extract failed: %v", err)
+	}
+
+	rc := getInflightRequests(t, ds, "m1")
+	if rc.Requests != 0 {
+		t.Errorf("expected Requests=0, got %d", rc.Requests)
+	}
+	if rc.Tokens != 0 {
+		t.Errorf("expected Tokens=0, got %d", rc.Tokens)
+	}
+}
+
+func TestCounterFloorsAtZero(t *testing.T) {
+	ext, ds := newInflightRequestsTest(t)
+
+	// Response with no prior request — both counters must floor at zero.
+	batch := []datalayer.Event{makeResponseEvent("m1", 50, 100)}
+	if err := ext.Extract(context.Background(), batch); err != nil {
+		t.Fatalf("Extract failed: %v", err)
+	}
+
+	rc := getInflightRequests(t, ds, "m1")
+	if rc.Requests != 0 {
+		t.Errorf("expected Requests=0, got %d", rc.Requests)
+	}
+	if rc.Tokens != 0 {
+		t.Errorf("expected Tokens=0, got %d", rc.Tokens)
+	}
+}
+
+func TestInflightRequestsMultipleModels(t *testing.T) {
+	ext, ds := newInflightRequestsTest(t)
+
+	batch := []datalayer.Event{
+		makeRequestEvent("m1", 10),
+		makeRequestEvent("m2", 20),
+	}
+	if err := ext.Extract(context.Background(), batch); err != nil {
+		t.Fatalf("Extract failed: %v", err)
+	}
+
+	rc1 := getInflightRequests(t, ds, "m1")
+	if rc1.Requests != 1 || rc1.Tokens != 10 {
+		t.Errorf("m1: expected {Requests:1, Tokens:10}, got %+v", rc1)
+	}
+
+	rc2 := getInflightRequests(t, ds, "m2")
+	if rc2.Requests != 1 || rc2.Tokens != 20 {
+		t.Errorf("m2: expected {Requests:1, Tokens:20}, got %+v", rc2)
+	}
+}
+
+func TestInflightRequestsUnknownEventTypeIgnored(t *testing.T) {
+	ext, ds := newInflightRequestsTest(t)
+
+	batch := []datalayer.Event{{Type: "unknown"}}
+	if err := ext.Extract(context.Background(), batch); err != nil {
+		t.Fatalf("Extract failed: %v", err)
+	}
+
+	ds.mu.Lock()
+	modelCount := len(ds.models)
+	ds.mu.Unlock()
+	if modelCount != 0 {
+		t.Errorf("expected no models in datastore, got %d", modelCount)
+	}
+}
+
+func TestInflightRequestsMissingModelFieldIgnored(t *testing.T) {
+	ext, ds := newInflightRequestsTest(t)
+
+	// Payload without a "model" key — no counter should be updated.
+	req := framework.NewInferenceRequest()
+	req.Body["max_tokens"] = float64(50)
+	batch := []datalayer.Event{
+		{Type: datalayer.RequestEventType, Payload: datalayer.RequestPayload{Request: req}},
+	}
+	if err := ext.Extract(context.Background(), batch); err != nil {
+		t.Fatalf("Extract failed: %v", err)
+	}
+
+	ds.mu.Lock()
+	modelCount := len(ds.models)
+	ds.mu.Unlock()
+	if modelCount != 0 {
+		t.Errorf("expected no models in datastore, got %d", modelCount)
+	}
+}