Request tracing support across core and connectors (#489)

* add OpenTelemetry tracing package for context propagation

* Propagate request context across event bus so subscribers inherit trace IDs

* Add trace_id, span_id and parent_span_id columns

* Include tracing for AMIE packet handlers

* Include tracing for COmanage subscriber

* Include tracing for SLURM subscribers

* source-column

* Add audit-driven trace store that reads audit_events and amie_audit_log

* expose audit endpoints

* Remove cross-connector tracing test from AMIE pipeline package

* Updated the baseline amie integration tests to reflect the additional tracing events

* Renumber tracing migrations

* Rename TraceEvent fields and COmanage subscriber audit event for clarity

* Add entity_type column to audit_events

* Replace amie_audit_log with amie_audit_extras table

* Route AMIE audit writes through core audit_events

* Drop UNION in audit_trace_store

* Add AMIE endpoints

* Update AMIE baseline integration tests to align with the new audit modeling

* Move AMIE HTTP endpoints into the connector via mux passthrough

* Consolidate AMIE migrations and switch trace/span IDs to hex strings

* Renumber audit migrations for consistency

Author: lahirujayathilake

README.md

@@ -47,6 +47,16 @@ airavata-custos/
 | `extensions/` | Independent services that run alongside Custos to extend HPC node behavior | `CILogon-SSH-PAM`, `SSH-Certificate-Signer` |
 | `dev-ops/` | Local dev stack and deployment automation | `compose/`, `terraform/`, `account-provisioning/` |
 
+## Audit conventions
+
+Every audit row in the system lives in the core `audit_events` table. Core, every connector, and any future extension write to it via the same shape: `event_type`, `entity_type`, `entity_id`, `details`, `source`, and the OpenTelemetry trace columns (`trace_id`, `span_id`, `parent_span_id`). The trace-view API at `/audit/traces*` reads from this one table.
+
+When a connector needs to attach connector-specific references to an audit row (for example AMIE keeps `packet_id` and `event_id` so it can fetch from a packet down to its audits), those references go in a separate `<connector>_audit_extras` table owned by the connector. The extras table has `audit_event_id` as its primary key with a `ON DELETE CASCADE` foreign key to `audit_events(id)`, plus whatever connector-specific columns it needs. The connector writes the row into its own extras table, inside the same transaction as the `audit_events` insert.
+
+Connector-specific endpoints live under `/connectors/{name}/...` and join `audit_events` with the connector's extras table. The unified trace view never reads the extras tables.
+
+The shape generalizes: a new connector that needs to record connector-specific references creates its own `<connector>_audit_extras` table and follows the same pattern. The core `audit_events` table stays neutral and never grows connector-shaped columns.
+
 ## Prerequisites
 
 * Go 1.24+

cmd/server/main.go

@@ -33,12 +33,14 @@ import (
 	"github.com/apache/airavata-custos/internal/connectors"
 	"github.com/apache/airavata-custos/internal/db"
 	"github.com/apache/airavata-custos/internal/server"
+	"github.com/apache/airavata-custos/internal/store"
+	"github.com/apache/airavata-custos/internal/tracing"
 	"github.com/apache/airavata-custos/pkg/events"
 	"github.com/apache/airavata-custos/pkg/service"
 )
 
 func main() {
-	slog.SetDefault(slog.New(slog.NewJSONHandler(os.Stdout, nil)))
+	slog.SetDefault(slog.New(tracing.SlogHandler(slog.NewJSONHandler(os.Stdout, nil))))
 
 	if err := run(); err != nil {
 		slog.Error("server exited with error", "error", err)
@@ -73,6 +75,26 @@ func run() error {
 		return err
 	}
 
+	tracingMode := tracing.ModeProduction
+	if os.Getenv("CUSTOS_TRACING_MODE") == "noop" {
+		tracingMode = tracing.ModeNoop
+	}
+	tracingShutdown, err := tracing.Init(tracing.InitConfig{
+		Mode:        tracingMode,
+		Logger:      slog.Default(),
+		ServiceName: "custos",
+	})
+	if err != nil {
+		return err
+	}
+	defer func() {
+		shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+		defer cancel()
+		if err := tracingShutdown(shutdownCtx); err != nil {
+			slog.Warn("tracing shutdown returned error", "error", err)
+		}
+	}()
+
 	// Create a new event bus instance to async messaging between service and connectors
 	eventBus := events.New()
 	svc := service.New(database, eventBus)
@@ -82,14 +104,19 @@ func run() error {
 
 	tryBootstrap(ctx, svc)
 
+	adminDeps := &server.AdminDeps{
+		AuditTraces: store.NewAuditTraceStore(database),
+	}
+	srv := server.New(svc, adminDeps)
+
 	// Tracks every background goroutine spawned by connectors so we can wait
 	// for them to drain on shutdown instead of killing them mid-flight.
 	var connectorsWG sync.WaitGroup
-	if err := connectors.LoadConnectors(ctx, database, eventBus, svc, &connectorsWG); err != nil {
+	if err := connectors.LoadConnectors(ctx, database, eventBus, svc, &connectorsWG, srv.Mux()); err != nil {
 		return err
 	}
 
-	handler := server.LoggingMiddleware(server.New(svc))
+	handler := server.LoggingMiddleware(tracing.Middleware(srv))
 
 	httpServer := &http.Server{
 		Addr:              addr,

connectors/ACCESS/AMIE-Processor/README.md

@@ -61,28 +61,93 @@ curl http://localhost:8083/metrics    # Prometheus metrics
 
 ## Testing
 
+Three layers, used for different things:
+
+### 1. Unit tests (correctness, no external services)
+
+```bash
+make test           # all tests, verbose
+make test-short     # short mode
+```
+
+Every package under this connector ships unit tests against `testify/mock`. No DB, no AMIE server, no network. Run on every commit.
+
+### 2. Integration tests (real DB + mock AMIE server)
+
+Run the per-handler suite (~11 handler tests) plus the pipeline suite (`baseline_integration_test`, `comanage_trace_chain_test`, `tracing_clean_slate_test`). All gated by build tag `integration` and require these env vars set:
+
+- `DATABASE_DSN`, `AMIE_BASE_URL`, `AMIE_SITE_CODE`, `AMIE_API_KEY`, `AMIE_CLUSTER_ID`
+
+The canonical runner brings up an isolated DB on `:3307` + mock AMIE on `:8181`, applies all migrations, fires every integration test, then tears down:
+
+```bash
+# From repo root
+make integration-test-amie
+```
+
+Equivalent script: `scripts/run-amie-integration-tests.sh`.
+
+#### What the pipeline tests fire
+
+Only one fixture exists today: `testdata/scenarios/baseline.yaml` — a deterministic 9-packet flow covering all major handler paths:
+
+| Packet | Asserts |
+|---|---|
+| `request_project_create` | PI user + project + allocation + PI cluster account |
+| `request_account_create` | user + cluster account + membership |
+| `data_project_create` / `data_account_create` | PERSIST_DNS rows |
+| `request_user_modify` | user + DN updates |
+| `request_person_merge` | survivor / retiree consolidation |
+| `request_account_inactivate` / `_reactivate` | membership status flips |
+| `request_project_inactivate` / `_reactivate` | project + all-member flips |
+| `inform_transaction_complete` | TRANSACTION_COMPLETE row |
+
+The three pipeline tests reuse the same fixture but assert different properties:
+- `baseline_integration_test` — DB shape + idempotency on rerun.
+- `comanage_trace_chain_test` — proves AMIE and COmanage audit rows share a `trace_id` over the unified `audit_events` table (uses mock COmanage REST; never hits the real registry).
+- `tracing_clean_slate_test` — clean-slate DB + full audit coverage + `/audit/*` endpoint smoke against an in-process httptest server.
+
+To add a new scenario: drop a YAML next to `baseline.yaml` and call `pipe.fireScenario(t, "<name>")` from a new `*_integration_test.go` file.
+
+### 3. Mock-server REST scenarios + k6 (load / soak / observability)
+
+For traffic generation against a **running** server (NOT a test — no assertions). Used to play in the admin UI, watch Grafana, or stress the worker.
+
+Start the mock AMIE server standalone:
+
 ```bash
-make test           # Run all tests with verbose output
-make test-short     # Run tests in short mode
+cd connectors/ACCESS/AMIE-Processor/mock-server
+source venv/bin/activate
+python3 mock-amie-server.py            # :8180
 ```
 
-All tests use mocks (testify/mock) and require no external services.
+Then queue packets via REST:
+
+```bash
+curl -X POST 'http://localhost:8180/test/TESTSITE/scenarios?type=mixed'         # success + failure mix
+curl -X POST 'http://localhost:8180/test/TESTSITE/scenarios?type=success_only'
+curl -X POST 'http://localhost:8180/test/TESTSITE/scenarios?type=failures_only'
+curl -X POST 'http://localhost:8180/test/TESTSITE/scenarios?type=heavy'          # large batch
+curl -X POST 'http://localhost:8180/test/TESTSITE/scenarios?type=dev_email'      # needs DEV_EMAIL env
+```
 
-**93 test functions, 160 total test cases (including subtests), 0 failures.**
+For sustained traffic at a configurable rate, use k6 with `mock-server/amie-traffic.js` — see `mock-server/README.md` for stage configuration.
 
 ## Observability
 
 The service exports Prometheus metrics at `/metrics`. A pre-built Grafana dashboard is available at `compose/grafana/dashboards/amie-service.json` showing packet processing stats, failures, retries, and processing duration percentiles.
 
-To run the full observability stack:
+To run the metrics stack:
 
 ```bash
 # From the repo root
-docker compose -f compose/docker-compose.yml up db prometheus grafana -d
+docker compose -f dev-ops/compose/docker-compose.yml up db prometheus grafana -d
 ```
 
 Then open Grafana at `http://localhost:3000` (admin/admin). The AMIE dashboard loads automatically.
 
+For request-flow tracing, every AMIE audit row carries `trace_id` / `span_id` / `parent_span_id` and lives in the core `audit_events` table (source=`amie`). The connector-specific `packet_id` / `event_id` references live in `amie_audit_extras` joined on `audit_event_id`. The admin trace view at `/audit/traces*` queries `audit_events` directly and renders the hierarchy from `parent_span_id`. For per-packet drill-down, see `GET /connectors/amie/packets/{packet_id}/audits`.
+
 ## Architecture
 
 ### Packet Processing Pipeline

connectors/ACCESS/AMIE-Processor/db/migrations/000001_initial_schema.down.sql

@@ -16,7 +16,7 @@
 -- under the License.
 
 DROP TABLE IF EXISTS amie_user_dns;
-DROP TABLE IF EXISTS amie_audit_log;
+DROP TABLE IF EXISTS amie_audit_extras;
 DROP TABLE IF EXISTS amie_processing_errors;
 DROP TABLE IF EXISTS amie_processing_events;
 DROP TABLE IF EXISTS amie_packets;

connectors/ACCESS/AMIE-Processor/db/migrations/000001_initial_schema.up.sql

@@ -74,22 +74,20 @@ CREATE TABLE IF NOT EXISTS amie_processing_errors
     KEY idx_amie_errors_occurred_at (occurred_at)
 ) ENGINE = InnoDB DEFAULT CHARSET = utf8mb4 COLLATE = utf8mb4_unicode_ci;
 
-CREATE TABLE IF NOT EXISTS amie_audit_log
+-- AMIE writes audit rows to core's audit_events table; this side table
+-- carries the connector-specific references (packet_id, event_id) joined on
+-- audit_event_id.
+CREATE TABLE IF NOT EXISTS amie_audit_extras
 (
-    id          BIGINT       NOT NULL AUTO_INCREMENT,
-    packet_id   VARCHAR(255) NOT NULL,
-    event_id    VARCHAR(255) NULL,
-    action      VARCHAR(64)  NOT NULL,
-    entity_type VARCHAR(64)  NOT NULL,
-    entity_id   VARCHAR(255) NULL,
-    summary     TEXT         NULL,
-    created_at  TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6),
-    PRIMARY KEY (id),
-    CONSTRAINT fk_amie_audit_packet FOREIGN KEY (packet_id) REFERENCES amie_packets (id) ON DELETE CASCADE,
-    CONSTRAINT fk_amie_audit_event FOREIGN KEY (event_id) REFERENCES amie_processing_events (id) ON DELETE SET NULL,
-    KEY idx_amie_audit_packet_id (packet_id),
-    KEY idx_amie_audit_action (action),
-    KEY idx_amie_audit_created_at (created_at)
+    audit_event_id VARCHAR(255) NOT NULL,
+    packet_id      VARCHAR(255) NOT NULL,
+    event_id       VARCHAR(255) NULL,
+    PRIMARY KEY (audit_event_id),
+    CONSTRAINT fk_amie_audit_extras_event  FOREIGN KEY (audit_event_id) REFERENCES audit_events(id) ON DELETE CASCADE,
+    CONSTRAINT fk_amie_audit_extras_packet FOREIGN KEY (packet_id) REFERENCES amie_packets(id) ON DELETE CASCADE,
+    CONSTRAINT fk_amie_audit_extras_procev FOREIGN KEY (event_id) REFERENCES amie_processing_events(id) ON DELETE SET NULL,
+    KEY idx_amie_audit_extras_packet_id (packet_id),
+    KEY idx_amie_audit_extras_event_id  (event_id)
 ) ENGINE = InnoDB DEFAULT CHARSET = utf8mb4 COLLATE = utf8mb4_unicode_ci;
 
 -- AMIE-side DN registry. AMIE delivers DnList fields that contain DNs across

connectors/ACCESS/AMIE-Processor/handler/data_account_create.go

@@ -24,8 +24,11 @@ import (
 	"fmt"
 	"log/slog"
 
+	"go.opentelemetry.io/otel/codes"
+
 	"github.com/apache/airavata-custos/connectors/ACCESS/AMIE-Processor/model"
 	"github.com/apache/airavata-custos/connectors/ACCESS/AMIE-Processor/store"
+	"github.com/apache/airavata-custos/internal/tracing"
 	"github.com/apache/airavata-custos/pkg/service"
 )
 
@@ -42,7 +45,16 @@ func NewDataAccountCreateHandler(svc *service.Service, userDNStore store.UserDNS
 
 func (h *DataAccountCreateHandler) SupportsType() string { return "data_account_create" }
 
-func (h *DataAccountCreateHandler) Handle(ctx context.Context, tx *sql.Tx, packetJSON map[string]any, packet *model.Packet, eventID string) error {
+func (h *DataAccountCreateHandler) Handle(ctx context.Context, tx *sql.Tx, packetJSON map[string]any, packet *model.Packet, eventID string) (err error) {
+	ctx, span := tracing.Start(ctx, "amie.handle:"+packet.Type)
+	defer span.End()
+	defer func() {
+		if err != nil {
+			span.RecordError(err)
+			span.SetStatus(codes.Error, err.Error())
+		}
+	}()
+
 	body, err := getBody(packetJSON)
 	if err != nil {
 		return err

connectors/ACCESS/AMIE-Processor/handler/data_project_create.go

@@ -24,8 +24,11 @@ import (
 	"fmt"
 	"log/slog"
 
+	"go.opentelemetry.io/otel/codes"
+
 	"github.com/apache/airavata-custos/connectors/ACCESS/AMIE-Processor/model"
 	"github.com/apache/airavata-custos/connectors/ACCESS/AMIE-Processor/store"
+	"github.com/apache/airavata-custos/internal/tracing"
 	"github.com/apache/airavata-custos/pkg/service"
 )
 
@@ -42,7 +45,16 @@ func NewDataProjectCreateHandler(svc *service.Service, userDNStore store.UserDNS
 
 func (h *DataProjectCreateHandler) SupportsType() string { return "data_project_create" }
 
-func (h *DataProjectCreateHandler) Handle(ctx context.Context, tx *sql.Tx, packetJSON map[string]any, packet *model.Packet, eventID string) error {
+func (h *DataProjectCreateHandler) Handle(ctx context.Context, tx *sql.Tx, packetJSON map[string]any, packet *model.Packet, eventID string) (err error) {
+	ctx, span := tracing.Start(ctx, "amie.handle:"+packet.Type)
+	defer span.End()
+	defer func() {
+		if err != nil {
+			span.RecordError(err)
+			span.SetStatus(codes.Error, err.Error())
+		}
+	}()
+
 	body, err := getBody(packetJSON)
 	if err != nil {
 		return err

connectors/ACCESS/AMIE-Processor/handler/handler.go

@@ -41,7 +41,8 @@ type AmieClient interface {
 	ReplyToPacket(ctx context.Context, packetRecID int64, reply map[string]any) error
 }
 
-// AuditService writes to amie_audit_log. The audit log is AMIE-local.
+// AuditService writes one audit_events row (source='amie') plus the matching
+// amie_audit_extras row carrying (packet_id, event_id).
 type AuditService interface {
 	Log(ctx context.Context, tx *sql.Tx, packetID, eventID string, action model.AuditAction, entityType, entityID, summary string) error
 }

connectors/ACCESS/AMIE-Processor/handler/inform_transaction_complete.go

@@ -23,7 +23,10 @@ import (
 	"fmt"
 	"log/slog"
 
+	"go.opentelemetry.io/otel/codes"
+
 	"github.com/apache/airavata-custos/connectors/ACCESS/AMIE-Processor/model"
+	"github.com/apache/airavata-custos/internal/tracing"
 )
 
 type InformTransactionCompleteHandler struct {
@@ -38,7 +41,16 @@ func (h *InformTransactionCompleteHandler) SupportsType() string {
 	return "inform_transaction_complete"
 }
 
-func (h *InformTransactionCompleteHandler) Handle(ctx context.Context, tx *sql.Tx, packetJSON map[string]any, packet *model.Packet, eventID string) error {
+func (h *InformTransactionCompleteHandler) Handle(ctx context.Context, tx *sql.Tx, packetJSON map[string]any, packet *model.Packet, eventID string) (err error) {
+	ctx, span := tracing.Start(ctx, "amie.handle:"+packet.Type)
+	defer span.End()
+	defer func() {
+		if err != nil {
+			span.RecordError(err)
+			span.SetStatus(codes.Error, err.Error())
+		}
+	}()
+
 	body, err := getBody(packetJSON)
 	if err != nil {
 		return err

connectors/ACCESS/AMIE-Processor/handler/integration_common.go

@@ -36,6 +36,7 @@ import (
 	amieservice "github.com/apache/airavata-custos/connectors/ACCESS/AMIE-Processor/service"
 	"github.com/apache/airavata-custos/connectors/ACCESS/AMIE-Processor/store"
 	"github.com/apache/airavata-custos/internal/db"
+	corestore "github.com/apache/airavata-custos/internal/store"
 	"github.com/apache/airavata-custos/pkg/events"
 	coreservice "github.com/apache/airavata-custos/pkg/service"
 )
@@ -102,7 +103,8 @@ func setupTestDB(t *testing.T) *sqlx.DB {
 func truncateAll(t *testing.T, database *sqlx.DB) {
 	t.Helper()
 	tables := []string{
-		"amie_audit_log",
+		"amie_audit_extras",
+		"audit_events",
 		"amie_processing_errors",
 		"amie_processing_events",
 		"amie_packets",
@@ -155,7 +157,7 @@ func newTestCoreService(database *sqlx.DB) *coreservice.Service {
 }
 
 func newTestAuditService(database *sqlx.DB) *amieservice.AuditService {
-	return amieservice.NewAuditService(store.NewAuditStore(database))
+	return amieservice.NewAuditService(corestore.NewAuditEventStore(database), store.NewAuditExtrasStore(database))
 }
 
 type fakeReply struct {
@@ -258,7 +260,9 @@ func countAuditActions(t *testing.T, database *sqlx.DB, packetID string, action
 	t.Helper()
 	var n int
 	if err := database.Get(&n,
-		"SELECT COUNT(*) FROM amie_audit_log WHERE packet_id = ? AND action = ?",
+		`SELECT COUNT(*) FROM audit_events ae
+		 JOIN amie_audit_extras x ON x.audit_event_id = ae.id
+		 WHERE x.packet_id = ? AND ae.event_type = ?`,
 		packetID, string(action),
 	); err != nil {
 		t.Fatalf("count audit %s: %v", action, err)