You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airavata.apache.org by ya...@apache.org on 2026/10/08 19:42:27 UTC
(airavata-custos) 02/02: Make trace ids mandatory and report trace status
This is an automated email from the ASF dual-hosted git repository.
yasithdev pushed a commit to branch amie-cr-fixes
in repository https://gitbox.apache.org/repos/asf/airavata-custos.git
commit eca27658952c5797c6b07777a8aad8f5bc7f2d15
Author: Yasith Jayawardana <ya...@icloud.com>
AuthorDate: Thu Oct 8 02:22:22 2026 -0400
Make trace ids mandatory and report trace status
---
api/core.openapi.yaml | 2 +
cmd/server/bootstrap.go | 3 +
cmd/server/main.go | 6 --
.../AMIE-Processor/pipeline/integration_common.go | 20 +----
.../ACCESS/AMIE-Processor/service/audit_service.go | 11 ---
.../AMIE-Processor/service/audit_service_test.go | 57 --------------
.../AMIE-Processor/store/packet_audit_store.go | 2 +-
connectors/ACCESS/AMIE-Processor/worker/poller.go | 3 +
.../internal/operations/mocks_test.go | 2 +-
.../internal/subscribers/reconciler.go | 3 +
.../Usage-Monitor/internal/smonitor/smonitor.go | 5 +-
.../internal/subscribers/subscriber.go | 3 +
internal/db/migrations/000005_audit.up.sql | 6 +-
internal/db/migrations/000007_event_bus.up.sql | 2 +-
internal/server/audit_handlers.go | 3 +-
internal/server/audit_handlers_test.go | 6 +-
internal/server/privilege_integration_test.go | 28 +++++++
internal/server/server.go | 2 +
internal/server/server_trace_test.go | 5 +-
internal/store/audit_event_store.go | 5 ++
internal/store/audit_trace_store.go | 37 ++++++---
.../store/audit_trace_store_integration_test.go | 70 +++++++++++++++++
internal/store/audit_trace_store_test.go | 9 ++-
.../store/event_delivery_store_integration_test.go | 1 +
internal/tracing/audit.go | 34 ---------
internal/tracing/middleware.go | 6 +-
internal/tracing/middleware_test.go | 34 ---------
internal/tracing/slog_handler.go | 10 +--
internal/tracing/tracing.go | 37 +++------
internal/tracing/tracing_test.go | 35 ++-------
pkg/events/integration_common_test.go | 6 --
pkg/events/worker.go | 18 ++---
pkg/models/audit.go | 2 +-
pkg/models/trace.go | 2 +-
pkg/service/audit_event.go | 9 ---
pkg/service/audit_event_test.go | 88 ----------------------
36 files changed, 208 insertions(+), 364 deletions(-)
diff --git a/api/core.openapi.yaml b/api/core.openapi.yaml
index a0b96e21c..2bb95fa68 100644
--- a/api/core.openapi.yaml
+++ b/api/core.openapi.yaml
@@ -1147,6 +1147,8 @@ paths:
description: OK
schema:
properties:
+ status:
+ type: string
trace_id:
type: string
tree:
diff --git a/cmd/server/bootstrap.go b/cmd/server/bootstrap.go
index 883288bbd..044a07539 100644
--- a/cmd/server/bootstrap.go
+++ b/cmd/server/bootstrap.go
@@ -22,6 +22,7 @@ import (
"log/slog"
"os"
+ "github.com/apache/airavata-custos/internal/tracing"
"github.com/apache/airavata-custos/pkg/service"
)
@@ -37,6 +38,8 @@ func tryBootstrap(ctx context.Context, svc *service.Service) {
if email == "" {
return
}
+ ctx, span := tracing.Start(ctx, "bootstrap.super_admin")
+ defer span.End()
if err := svc.BootstrapSuperAdmin(ctx, email, "env:"+bootstrapAdminEmailEnv); err != nil {
slog.Warn("bootstrap super_admin failed", "email", email, "error", err)
}
diff --git a/cmd/server/main.go b/cmd/server/main.go
index 375faf5a8..882998f0a 100644
--- a/cmd/server/main.go
+++ b/cmd/server/main.go
@@ -102,13 +102,7 @@ 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 {
diff --git a/connectors/ACCESS/AMIE-Processor/pipeline/integration_common.go b/connectors/ACCESS/AMIE-Processor/pipeline/integration_common.go
index 84699d604..5f3215d24 100644
--- a/connectors/ACCESS/AMIE-Processor/pipeline/integration_common.go
+++ b/connectors/ACCESS/AMIE-Processor/pipeline/integration_common.go
@@ -22,7 +22,6 @@ package pipeline
import (
"context"
"fmt"
- "log/slog"
"net/http"
"os"
"strings"
@@ -44,7 +43,6 @@ import (
"github.com/apache/airavata-custos/connectors/ACCESS/AMIE-Processor/worker"
"github.com/apache/airavata-custos/internal/db"
corestore "github.com/apache/airavata-custos/internal/store"
- "github.com/apache/airavata-custos/internal/tracing"
"github.com/apache/airavata-custos/pkg/events"
coreservice "github.com/apache/airavata-custos/pkg/service"
)
@@ -52,11 +50,9 @@ import (
const testClusterID = "00000000-0000-0000-0000-000000000001"
var (
- sharedDB *sqlx.DB
- sharedDBOnce sync.Once
- sharedDBErr error
- tracingInitOnce sync.Once
- tracingInitError error
+ sharedDB *sqlx.DB
+ sharedDBOnce sync.Once
+ sharedDBErr error
)
func isLocalAMIEConfigAvailable() bool {
@@ -102,16 +98,6 @@ func setupTestDB(t *testing.T) *sqlx.DB {
if sharedDBErr != nil {
t.Fatalf("setup db: %v", sharedDBErr)
}
- tracingInitOnce.Do(func() {
- _, tracingInitError = tracing.Init(tracing.InitConfig{
- Mode: tracing.ModeProduction,
- Logger: slog.Default(),
- ServiceName: "custos",
- })
- })
- if tracingInitError != nil {
- t.Fatalf("tracing init: %v", tracingInitError)
- }
truncateAll(t, sharedDB)
seedCluster(t, sharedDB)
return sharedDB
diff --git a/connectors/ACCESS/AMIE-Processor/service/audit_service.go b/connectors/ACCESS/AMIE-Processor/service/audit_service.go
index ebacf7c70..e46e76fa4 100644
--- a/connectors/ACCESS/AMIE-Processor/service/audit_service.go
+++ b/connectors/ACCESS/AMIE-Processor/service/audit_service.go
@@ -28,7 +28,6 @@ import (
"github.com/apache/airavata-custos/connectors/ACCESS/AMIE-Processor/model"
"github.com/apache/airavata-custos/internal/store"
- "github.com/apache/airavata-custos/internal/tracing"
"github.com/apache/airavata-custos/pkg/models"
)
@@ -69,16 +68,6 @@ func (s *AuditService) Log(ctx context.Context, tx *sql.Tx, packetID, eventID st
Details: summary,
Source: auditSource,
}
- tracing.PopulateAuditIDs(ctx, &event.TraceID, &event.SpanID, &event.ParentSpanID)
- if event.TraceID == "" {
- slog.WarnContext(ctx, "audit write outside an active span",
- "packet_id", packetID,
- "event_id", eventID,
- "action", string(action),
- "entity_type", entityType,
- "entity_id", entityID,
- )
- }
if err := s.coreEvents.Create(ctx, tx, event); err != nil {
return fmt.Errorf("audit_service: saving audit_events row for packet %s: %w", packetID, err)
diff --git a/connectors/ACCESS/AMIE-Processor/service/audit_service_test.go b/connectors/ACCESS/AMIE-Processor/service/audit_service_test.go
index b90ba091c..d5f820856 100644
--- a/connectors/ACCESS/AMIE-Processor/service/audit_service_test.go
+++ b/connectors/ACCESS/AMIE-Processor/service/audit_service_test.go
@@ -25,12 +25,8 @@ import (
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
- "go.opentelemetry.io/otel"
- sdktrace "go.opentelemetry.io/otel/sdk/trace"
- "go.opentelemetry.io/otel/trace"
"github.com/apache/airavata-custos/connectors/ACCESS/AMIE-Processor/model"
- "github.com/apache/airavata-custos/internal/tracing"
"github.com/apache/airavata-custos/pkg/models"
)
@@ -154,59 +150,6 @@ func TestAuditLog_PropagatesExtrasStoreError(t *testing.T) {
assert.Contains(t, err.Error(), "amie_audit_extras")
}
-func TestAuditLog_PersistsTraceAndSpanIDs(t *testing.T) {
- prev := otel.GetTracerProvider()
- tp := sdktrace.NewTracerProvider()
- otel.SetTracerProvider(tp)
- t.Cleanup(func() { otel.SetTracerProvider(prev) })
-
- ctx, span := tracing.Start(context.Background(), "test.root")
- defer span.End()
- wantTrace := span.SpanContext().TraceID()
- wantSpan := span.SpanContext().SpanID()
-
- coreEvents := new(mockCoreEventStore)
- extras := new(mockExtrasStore)
- svc := NewAuditService(coreEvents, extras)
-
- var captured *models.AuditEvent
- coreEvents.On("Create", mock.Anything, mock.Anything, mock.MatchedBy(func(e *models.AuditEvent) bool {
- captured = e
- return true
- })).Return(nil)
- extras.On("Save", mock.Anything, mock.Anything, mock.Anything).Return(nil)
-
- require.NoError(t, svc.Log(ctx, nil, "packet-trace", "event-trace", model.AuditCreatePerson, "person", "p1", "with trace"))
- require.NotNil(t, captured)
- if captured.TraceID != wantTrace.String() {
- t.Fatalf("trace_id mismatch: got %s want %s", captured.TraceID, wantTrace.String())
- }
- if captured.SpanID != wantSpan.String() {
- t.Fatalf("span_id mismatch: got %s want %s", captured.SpanID, wantSpan.String())
- }
-}
-
-func TestAuditLog_NilTraceWhenNoSpan(t *testing.T) {
- coreEvents := new(mockCoreEventStore)
- extras := new(mockExtrasStore)
- svc := NewAuditService(coreEvents, extras)
-
- var captured *models.AuditEvent
- coreEvents.On("Create", mock.Anything, mock.Anything, mock.MatchedBy(func(e *models.AuditEvent) bool {
- captured = e
- return true
- })).Return(nil)
- extras.On("Save", mock.Anything, mock.Anything, mock.Anything).Return(nil)
-
- require.NoError(t, svc.Log(
- trace.ContextWithSpanContext(context.Background(), trace.SpanContext{}),
- nil, "p", "e", model.AuditReplySent, "reply", "", ""))
- require.NotNil(t, captured)
- if captured.TraceID != "" || captured.SpanID != "" {
- t.Fatalf("expected nil trace/span IDs when no active span")
- }
-}
-
func TestAuditLog_RejectsEmptyPacketID(t *testing.T) {
svc := NewAuditService(new(mockCoreEventStore), new(mockExtrasStore))
err := svc.Log(context.Background(), nil, "", "evt", model.AuditCreatePerson, "person", "p", "")
diff --git a/connectors/ACCESS/AMIE-Processor/store/packet_audit_store.go b/connectors/ACCESS/AMIE-Processor/store/packet_audit_store.go
index 20d7d7029..c939a66a3 100644
--- a/connectors/ACCESS/AMIE-Processor/store/packet_audit_store.go
+++ b/connectors/ACCESS/AMIE-Processor/store/packet_audit_store.go
@@ -61,7 +61,7 @@ ORDER BY ae.event_time ASC, ae.span_id ASC
func (s *pgPacketAuditStore) ListAuditsForPacket(ctx context.Context, packetID string) ([]models.TraceEvent, error) {
type row struct {
SpanID string `db:"span_id"`
- ParentSpanID string `db:"parent_span_id"`
+ ParentSpanID *string `db:"parent_span_id"`
Source string `db:"source"`
EventType string `db:"event_type"`
EntityType string `db:"entity_type"`
diff --git a/connectors/ACCESS/AMIE-Processor/worker/poller.go b/connectors/ACCESS/AMIE-Processor/worker/poller.go
index 83a08ab68..58a5862d8 100644
--- a/connectors/ACCESS/AMIE-Processor/worker/poller.go
+++ b/connectors/ACCESS/AMIE-Processor/worker/poller.go
@@ -31,6 +31,7 @@ import (
"github.com/apache/airavata-custos/connectors/ACCESS/AMIE-Processor/amieclient"
custosdb "github.com/apache/airavata-custos/connectors/ACCESS/AMIE-Processor/db"
"github.com/apache/airavata-custos/connectors/ACCESS/AMIE-Processor/model"
+ "github.com/apache/airavata-custos/internal/tracing"
)
type pollerPacketStore interface {
@@ -98,6 +99,8 @@ func (p *Poller) Run(ctx context.Context) {
}
func (p *Poller) pollForPackets(ctx context.Context) {
+ ctx, span := tracing.Start(ctx, "amie.poll")
+ defer span.End()
slog.Info("Polling for new AMIE packets...")
packets, err := p.client.FetchInProgressPackets(ctx)
if err != nil {
diff --git a/connectors/COmanage/Identity-Provisioner/internal/operations/mocks_test.go b/connectors/COmanage/Identity-Provisioner/internal/operations/mocks_test.go
index 7bab3fb56..76a88fed3 100644
--- a/connectors/COmanage/Identity-Provisioner/internal/operations/mocks_test.go
+++ b/connectors/COmanage/Identity-Provisioner/internal/operations/mocks_test.go
@@ -81,7 +81,7 @@ func (f *fakeCore) ListMembersForAllocation(_ context.Context, _ string) ([]stor
}
func (f *fakeCore) CreateAuditEvent(ctx context.Context, e *models.AuditEvent) (*models.AuditEvent, error) {
- tracing.PopulateAuditIDs(ctx, &e.TraceID, &e.SpanID, &e.ParentSpanID)
+ e.TraceID, e.SpanID = tracing.IDsFromContext(ctx)
if e.ID == "" {
e.ID = fmt.Sprintf("audit-%d", len(f.auditEvents)+1)
}
diff --git a/connectors/SLURM/Association-Mapper/internal/subscribers/reconciler.go b/connectors/SLURM/Association-Mapper/internal/subscribers/reconciler.go
index c761a4bac..695878a89 100644
--- a/connectors/SLURM/Association-Mapper/internal/subscribers/reconciler.go
+++ b/connectors/SLURM/Association-Mapper/internal/subscribers/reconciler.go
@@ -23,6 +23,7 @@ import (
"time"
"github.com/apache/airavata-custos/connectors/SLURM/Rest-Client/pkg/client"
+ "github.com/apache/airavata-custos/internal/tracing"
"github.com/apache/airavata-custos/pkg/models"
)
@@ -75,6 +76,8 @@ func (a *AssociationSubscriber) StartReconciler(ctx context.Context) {
// access, but a lost deactivation leaves a member able to submit indefinitely
// with nothing reporting it.
func (a *AssociationSubscriber) reconcile(ctx context.Context) {
+ ctx, span := tracing.Start(ctx, "slurm.reconcile")
+ defer span.End()
sweepCtx, cancel := context.WithTimeout(ctx, reconcileSweepTimeout)
defer cancel()
diff --git a/connectors/SLURM/Usage-Monitor/internal/smonitor/smonitor.go b/connectors/SLURM/Usage-Monitor/internal/smonitor/smonitor.go
index 9aaf30fdb..6e7ca3d5c 100644
--- a/connectors/SLURM/Usage-Monitor/internal/smonitor/smonitor.go
+++ b/connectors/SLURM/Usage-Monitor/internal/smonitor/smonitor.go
@@ -24,6 +24,7 @@ import (
"time"
"github.com/apache/airavata-custos/connectors/SLURM/Rest-Client/pkg/client"
+ "github.com/apache/airavata-custos/internal/tracing"
"github.com/apache/airavata-custos/pkg/events"
"github.com/apache/airavata-custos/pkg/models"
"github.com/apache/airavata-custos/pkg/service"
@@ -105,7 +106,9 @@ func (m *SlurmMonitor) StartMonitor(ctx context.Context) {
func (m *SlurmMonitor) poll() {
slog.Debug("polling SLURM usage")
- context, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+ ctx, span := tracing.Start(context.Background(), "slurm.usage_poll")
+ defer span.End()
+ context, cancel := context.WithTimeout(ctx, 10*time.Second)
defer cancel()
cluster, err := m.coreService.GetComputeCluster(context, m.clusterId)
if err != nil {
diff --git a/connectors/Storage/VAST-Provisioner/internal/subscribers/subscriber.go b/connectors/Storage/VAST-Provisioner/internal/subscribers/subscriber.go
index 407a1f8a9..56461f318 100644
--- a/connectors/Storage/VAST-Provisioner/internal/subscribers/subscriber.go
+++ b/connectors/Storage/VAST-Provisioner/internal/subscribers/subscriber.go
@@ -29,6 +29,7 @@ import (
"time"
"github.com/apache/airavata-custos/connectors/Storage/VAST-Provisioner/internal/client"
+ "github.com/apache/airavata-custos/internal/tracing"
"github.com/apache/airavata-custos/pkg/events"
"github.com/apache/airavata-custos/pkg/models"
"github.com/apache/airavata-custos/pkg/service"
@@ -166,6 +167,8 @@ func (s *StorageSubscriber) StartReconciler(ctx context.Context) {
}
func (s *StorageSubscriber) reconcile(ctx context.Context, id string, c Cluster) {
+ ctx, span := tracing.Start(ctx, "vast.reconcile")
+ defer span.End()
users, err := s.core.ListComputeClusterUsersByCluster(ctx, id)
errs := []error{err}
for _, cu := range users {
diff --git a/internal/db/migrations/000005_audit.up.sql b/internal/db/migrations/000005_audit.up.sql
index e315202b2..f057281d8 100644
--- a/internal/db/migrations/000005_audit.up.sql
+++ b/internal/db/migrations/000005_audit.up.sql
@@ -27,9 +27,9 @@ CREATE TABLE IF NOT EXISTS audit_events
entity_type VARCHAR(64) NOT NULL DEFAULT '',
details TEXT NOT NULL,
source VARCHAR(64) NOT NULL DEFAULT 'core',
- trace_id CHAR(32) NOT NULL DEFAULT '',
- span_id CHAR(16) NOT NULL DEFAULT '',
- parent_span_id CHAR(16) NOT NULL DEFAULT '',
+ trace_id CHAR(32) NOT NULL CHECK (trace_id <> ''),
+ span_id CHAR(16) NOT NULL CHECK (span_id <> ''),
+ parent_span_id CHAR(16) CHECK (parent_span_id <> ''),
PRIMARY KEY (id)
);
CREATE INDEX IF NOT EXISTS idx_audit_events_entity ON audit_events (entity_id, event_time);
diff --git a/internal/db/migrations/000007_event_bus.up.sql b/internal/db/migrations/000007_event_bus.up.sql
index 53b22184a..58a4fd41d 100644
--- a/internal/db/migrations/000007_event_bus.up.sql
+++ b/internal/db/migrations/000007_event_bus.up.sql
@@ -35,7 +35,7 @@ CREATE TABLE IF NOT EXISTS events
event_type VARCHAR(128) NOT NULL,
payload JSONB NOT NULL,
source VARCHAR(64) NOT NULL,
- trace_id CHAR(32) NOT NULL DEFAULT '',
+ trace_id CHAR(32) NOT NULL CHECK (trace_id <> ''),
created_at TIMESTAMPTZ(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6),
PRIMARY KEY (id)
);
diff --git a/internal/server/audit_handlers.go b/internal/server/audit_handlers.go
index 7307df47a..d9c55a655 100644
--- a/internal/server/audit_handlers.go
+++ b/internal/server/audit_handlers.go
@@ -95,7 +95,7 @@ func (s *Server) handleListTraces(w http.ResponseWriter, r *http.Request) {
// @Security BearerAuth
// @Produce json
// @Param trace_id path string true "32-char lowercase hex"
-// @Success 200 {object} object{trace_id=string,tree=[]models.TraceNode,truncated=boolean}
+// @Success 200 {object} object{trace_id=string,status=string,tree=[]models.TraceNode,truncated=boolean}
// @Failure 400 {object} object{error=string} "Malformed trace_id"
// @Failure 404 {object} object{error=string} "Trace not found"
// @Failure 503 {object} object{error=string} "Audit trace store not configured"
@@ -122,6 +122,7 @@ func (s *Server) handleGetTrace(w http.ResponseWriter, r *http.Request) {
resp := map[string]any{
"trace_id": traceID,
+ "status": tree.Status,
"tree": tree.Children,
"truncated": truncated,
}
diff --git a/internal/server/audit_handlers_test.go b/internal/server/audit_handlers_test.go
index ec08fd207..20bd2fb70 100644
--- a/internal/server/audit_handlers_test.go
+++ b/internal/server/audit_handlers_test.go
@@ -246,6 +246,7 @@ func TestGetTraceReturnsTree(t *testing.T) {
rootSpan := "0102030405060708"
childSpan := "090a0b0c0d0e0f10"
tree := &models.TraceNode{
+ TraceEvent: models.TraceEvent{Status: "error"},
Children: []*models.TraceNode{
{
TraceEvent: models.TraceEvent{
@@ -259,7 +260,7 @@ func TestGetTraceReturnsTree(t *testing.T) {
{
TraceEvent: models.TraceEvent{
SpanID: childSpan,
- ParentSpanID: rootSpan,
+ ParentSpanID: &rootSpan,
Source: "comanage",
EventType: "ComanageProvisioningFailed",
Status: "error",
@@ -288,6 +289,9 @@ func TestGetTraceReturnsTree(t *testing.T) {
if body["trace_id"].(string) != raw {
t.Errorf("trace_id = %v", body["trace_id"])
}
+ if body["status"] != "error" {
+ t.Errorf("trace status = %v", body["status"])
+ }
treeNodes := body["tree"].([]any)
if len(treeNodes) != 1 {
t.Fatalf("tree len = %d", len(treeNodes))
diff --git a/internal/server/privilege_integration_test.go b/internal/server/privilege_integration_test.go
index 87f4d5d85..85c01c05a 100644
--- a/internal/server/privilege_integration_test.go
+++ b/internal/server/privilege_integration_test.go
@@ -224,6 +224,34 @@ func TestGrantPrivilegeEndpoint_HappyPath(t *testing.T) {
}
}
+// Make sure a request through the server chain stamps its audit row with the
+// request's trace, the one returned in X-Trace-Id.
+func TestGrantPrivilegeEndpoint_AuditRowCarriesRequestTrace(t *testing.T) {
+ database, _, srv := setupTestStack(t)
+ granter := seedUser(t, database, "granter@example.edu")
+ target := seedUser(t, database, "target@example.edu")
+ seedPrivilegesGrant(t, database, granter)
+
+ body, _ := json.Marshal(map[string]any{"privilege": "core:clusters:read"})
+ rr := httptest.NewRecorder()
+ req := httptest.NewRequest(http.MethodPost, "/users/"+target+"/privileges", bytes.NewReader(body))
+ LoggingMiddleware(srv).ServeHTTP(rr, withTestCaller(req, granter, models.PrivilegesGrant))
+ if rr.Code != http.StatusCreated {
+ t.Fatalf("status: got %d, want 201, body=%s", rr.Code, rr.Body.String())
+ }
+
+ var row struct {
+ TraceID string `db:"trace_id"`
+ SpanID string `db:"span_id"`
+ }
+ if err := database.Get(&row, "SELECT trace_id, span_id FROM audit_events WHERE entity_id = $1", target); err != nil {
+ t.Fatalf("read audit row: %v", err)
+ }
+ if want := rr.Header().Get("X-Trace-Id"); want == "" || row.TraceID != want || row.SpanID == "" {
+ t.Fatalf("audit row trace=%q span=%q, want the request's trace %q", row.TraceID, row.SpanID, want)
+ }
+}
+
func TestGrantPrivilegeEndpoint_GranterWithoutMeta_403(t *testing.T) {
database, _, srv := setupTestStack(t)
plain := seedUser(t, database, "plain@example.edu")
diff --git a/internal/server/server.go b/internal/server/server.go
index 9cf78be75..cca8eace1 100644
--- a/internal/server/server.go
+++ b/internal/server/server.go
@@ -26,6 +26,7 @@ import (
"time"
"github.com/apache/airavata-custos/internal/httputil"
+ "github.com/apache/airavata-custos/internal/tracing"
"github.com/apache/airavata-custos/pkg/common"
"github.com/apache/airavata-custos/pkg/identity"
"github.com/apache/airavata-custos/pkg/models"
@@ -215,6 +216,7 @@ type statusUpdateRequest struct {
// from the X-Trace-Id response header the inner span set, since the inner
// request ctx is not visible at this scope after ServeHTTP returns.
func LoggingMiddleware(next http.Handler) http.Handler {
+ next = tracing.Middleware(next)
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
start := time.Now()
rw := &httputil.StatusRecorder{ResponseWriter: w, Status: http.StatusOK}
diff --git a/internal/server/server_trace_test.go b/internal/server/server_trace_test.go
index 40f86525c..b15fcc323 100644
--- a/internal/server/server_trace_test.go
+++ b/internal/server/server_trace_test.go
@@ -26,12 +26,11 @@ import (
sdktrace "go.opentelemetry.io/otel/sdk/trace"
"go.opentelemetry.io/otel/sdk/trace/tracetest"
- "github.com/apache/airavata-custos/internal/tracing"
"github.com/apache/airavata-custos/pkg/identity"
)
// TestLoggingWrapsTracingProducesTraceIdHeader verifies the cmd/server stack
-// composition: LoggingMiddleware(tracing.Middleware(server.New(...))) lets the
+// composition: LoggingMiddleware(server.New(...)) lets the
// access log see the trace_id and surfaces X-Trace-Id to the client.
func TestLoggingWrapsTracingProducesTraceIdHeader(t *testing.T) {
prev := otel.GetTracerProvider()
@@ -41,7 +40,7 @@ func TestLoggingWrapsTracingProducesTraceIdHeader(t *testing.T) {
t.Cleanup(func() { otel.SetTracerProvider(prev) })
router := identity.NewRouter(http.NewServeMux())
- handler := LoggingMiddleware(tracing.Middleware(New(nil, router)))
+ handler := LoggingMiddleware(New(nil, router))
req := httptest.NewRequest(http.MethodGet, "/healthz", nil)
rec := httptest.NewRecorder()
diff --git a/internal/store/audit_event_store.go b/internal/store/audit_event_store.go
index fb849dd87..fd6503b07 100644
--- a/internal/store/audit_event_store.go
+++ b/internal/store/audit_event_store.go
@@ -24,6 +24,7 @@ import (
"github.com/jmoiron/sqlx"
+ "github.com/apache/airavata-custos/internal/tracing"
"github.com/apache/airavata-custos/pkg/models"
)
@@ -90,6 +91,10 @@ func (s *pgAuditEventStore) ListAll(ctx context.Context) ([]*models.AuditEvent,
}
func (s *pgAuditEventStore) Create(ctx context.Context, tx *sql.Tx, e *models.AuditEvent) error {
+ if e.TraceID == "" {
+ e.TraceID, e.SpanID = tracing.IDsFromContext(ctx)
+ e.ParentSpanID = tracing.ParentSpanIDFromContext(ctx)
+ }
_, err := tx.ExecContext(ctx,
`INSERT INTO audit_events (id, event_type, event_time, entity_id, entity_type, details, source, trace_id, span_id, parent_span_id)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)`,
diff --git a/internal/store/audit_trace_store.go b/internal/store/audit_trace_store.go
index 6c0159a3f..967f6785e 100644
--- a/internal/store/audit_trace_store.go
+++ b/internal/store/audit_trace_store.go
@@ -66,7 +66,6 @@ SELECT
details AS description,
event_time AS created_at
FROM audit_events
-WHERE trace_id <> ''
`
type pgAuditTraceStore struct {
@@ -80,7 +79,7 @@ func NewAuditTraceStore(db *sqlx.DB) AuditTraceStore {
type rowEvent struct {
TraceID string `db:"trace_id"`
SpanID string `db:"span_id"`
- ParentSpanID string `db:"parent_span_id"`
+ ParentSpanID *string `db:"parent_span_id"`
Source string `db:"source"`
EventType string `db:"event_type"`
EntityType string `db:"entity_type"`
@@ -107,6 +106,24 @@ func (r rowEvent) toTraceEvent() models.TraceEvent {
// Acceptable while page sizes stay <= 200.
func (s *pgAuditTraceStore) ListTraces(ctx context.Context, f TraceFilter) ([]models.TraceSummary, int, error) {
whereSQL, args := buildTraceWhere(f)
+ // Status is computed in Go, so narrow to matching traces before counting and paging.
+ if len(f.Statuses) > 0 {
+ var ids, kept []string
+ if err := s.db.SelectContext(ctx, &ids, s.db.Rebind(`SELECT DISTINCT trace_id FROM (`+auditTraceSelect+`) u `+whereSQL), args...); err != nil {
+ return nil, 0, fmt.Errorf("audit_trace_store: candidates: %w", err)
+ }
+ for _, id := range ids {
+ _, _, status, err := s.summariseTrace(ctx, id)
+ if err != nil {
+ return nil, 0, err
+ }
+ if matchesStatusFilter(status, f.Statuses) {
+ kept = append(kept, id)
+ }
+ }
+ whereSQL = "WHERE u.trace_id = ANY(?)" + strings.Replace(whereSQL, "WHERE", " AND", 1)
+ args = append([]any{kept}, args...)
+ }
limit := f.Limit
if limit <= 0 {
@@ -150,9 +167,6 @@ func (s *pgAuditTraceStore) ListTraces(ctx context.Context, f TraceFilter) ([]mo
if err != nil {
return nil, 0, err
}
- if !matchesStatusFilter(status, f.Statuses) {
- continue
- }
out = append(out, models.TraceSummary{
TraceID: r.TraceID,
RootOperation: rootOp,
@@ -181,7 +195,7 @@ func (s *pgAuditTraceStore) summariseTrace(ctx context.Context, traceID string)
var rootOp string
for _, r := range rows {
- if r.ParentSpanID == "" {
+ if r.ParentSpanID == nil {
rootOp = r.EventType
break
}
@@ -200,7 +214,7 @@ func (s *pgAuditTraceStore) summariseTrace(ctx context.Context, traceID string)
func dominantSource(rows []rowEvent) string {
for _, r := range rows {
- if r.ParentSpanID == "" {
+ if r.ParentSpanID == nil {
return r.Source
}
}
@@ -225,7 +239,10 @@ func (s *pgAuditTraceStore) GetTraceTree(ctx context.Context, traceID string) (*
rows = rows[:TreeRowLimit]
}
- return buildTree(rows), truncated, nil
+ _, _, status, err := s.summariseTrace(ctx, traceID)
+ tree := buildTree(rows)
+ tree.Status = status
+ return tree, truncated, err
}
// buildTree links rows by span_id. Rows whose parent isn't in the table
@@ -238,8 +255,8 @@ func buildTree(rows []rowEvent) *models.TraceNode {
root := &models.TraceNode{}
for _, r := range rows {
node := bySpan[r.SpanID]
- if r.ParentSpanID != "" {
- if parent, ok := bySpan[r.ParentSpanID]; ok {
+ if r.ParentSpanID != nil {
+ if parent, ok := bySpan[*r.ParentSpanID]; ok {
parent.Children = append(parent.Children, node)
continue
}
diff --git a/internal/store/audit_trace_store_integration_test.go b/internal/store/audit_trace_store_integration_test.go
new file mode 100644
index 000000000..409c2041b
--- /dev/null
+++ b/internal/store/audit_trace_store_integration_test.go
@@ -0,0 +1,70 @@
+//go:build integration
+
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you 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 store
+
+import (
+ "context"
+ "fmt"
+ "testing"
+ "time"
+)
+
+// Make sure a status filter narrows the trace set before counting and paging,
+// and that null parent span ids read back as roots.
+func TestListTracesStatusFilterCountsAndPages(t *testing.T) {
+ database := setupTestDB(t)
+ ctx := context.Background()
+ const source = "trace-status-test"
+ if _, err := database.Exec(`DELETE FROM audit_events WHERE source = $1`, source); err != nil {
+ t.Fatalf("clean: %v", err)
+ }
+ base := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)
+ // Traces alternate error / in_progress; the newest trace is index 5.
+ for i := range 6 {
+ traceID := fmt.Sprintf("%032x", 0xa00+i)
+ eventType := "ROOT_STARTED"
+ if i%2 == 0 {
+ eventType = "ROOT_FAILED"
+ }
+ if _, err := database.Exec(`INSERT INTO audit_events
+ (id, event_type, event_time, entity_id, details, source, trace_id, span_id, parent_span_id)
+ VALUES ($1, $2, $3, 'e', '', $4, $5, $6, NULL)`,
+ traceID, eventType, base.Add(time.Duration(i)*time.Minute), source, traceID, traceID[:16]); err != nil {
+ t.Fatalf("insert: %v", err)
+ }
+ }
+
+ s := NewAuditTraceStore(database)
+ got, total, err := s.ListTraces(ctx, TraceFilter{Sources: []string{source}, Statuses: []string{"error"}, Limit: 2, Offset: 2})
+ if err != nil {
+ t.Fatalf("list: %v", err)
+ }
+ if total != 3 || len(got) != 1 || got[0].TraceID != fmt.Sprintf("%032x", 0xa00) || got[0].Status != "error" {
+ t.Fatalf("total=%d page=%+v, want total 3 and the oldest error trace", total, got)
+ }
+
+ tree, _, err := s.GetTraceTree(ctx, got[0].TraceID)
+ if err != nil {
+ t.Fatalf("tree: %v", err)
+ }
+ if tree.Status != "error" || len(tree.Children) != 1 || tree.Children[0].ParentSpanID != nil {
+ t.Fatalf("tree status=%q children=%+v, want error with one root", tree.Status, tree.Children)
+ }
+}
diff --git a/internal/store/audit_trace_store_test.go b/internal/store/audit_trace_store_test.go
index fb753c0dd..d428e205e 100644
--- a/internal/store/audit_trace_store_test.go
+++ b/internal/store/audit_trace_store_test.go
@@ -35,8 +35,8 @@ func TestBuildTreeNestsByParent(t *testing.T) {
leaf := mkSpan(0x03)
rows := []rowEvent{
{SpanID: root, Source: "amie", EventType: "CREATE_PERSON", CreatedAt: time.Unix(1, 0)},
- {SpanID: child, ParentSpanID: root, Source: "comanage", EventType: "ComanageLookup", CreatedAt: time.Unix(2, 0)},
- {SpanID: leaf, ParentSpanID: child, Source: "comanage", EventType: "ComanageClusterAccountAttached", CreatedAt: time.Unix(3, 0)},
+ {SpanID: child, ParentSpanID: &root, Source: "comanage", EventType: "ComanageLookup", CreatedAt: time.Unix(2, 0)},
+ {SpanID: leaf, ParentSpanID: &child, Source: "comanage", EventType: "ComanageClusterAccountAttached", CreatedAt: time.Unix(3, 0)},
}
tree := buildTree(rows)
if got := len(tree.Children); got != 1 {
@@ -64,7 +64,7 @@ func TestBuildTreeOrphansBecomeTopLevel(t *testing.T) {
rows := []rowEvent{
{SpanID: row1, Source: "amie", EventType: "ROOT"},
// parent references a span that did not write an audit row
- {SpanID: row2, ParentSpanID: ghost, Source: "amie", EventType: "ORPHAN"},
+ {SpanID: row2, ParentSpanID: &ghost, Source: "amie", EventType: "ORPHAN"},
}
tree := buildTree(rows)
if got := len(tree.Children); got != 2 {
@@ -132,8 +132,9 @@ func TestMatchesStatusFilter(t *testing.T) {
}
func TestDominantSourcePicksRoot(t *testing.T) {
+ parent := mkSpan(0x01)
rows := []rowEvent{
- {Source: "comanage", ParentSpanID: mkSpan(0x01)},
+ {Source: "comanage", ParentSpanID: &parent},
{Source: "amie"}, // root: nil parent
}
if got := dominantSource(rows); got != "amie" {
diff --git a/internal/store/event_delivery_store_integration_test.go b/internal/store/event_delivery_store_integration_test.go
index 92e22dfcf..665b8a40e 100644
--- a/internal/store/event_delivery_store_integration_test.go
+++ b/internal/store/event_delivery_store_integration_test.go
@@ -38,6 +38,7 @@ func newEvent(createdAt time.Time) *models.Event {
EventType: "compute_cluster_user::create",
Payload: json.RawMessage(`{"id":"cu-1","local_username":"jdoe"}`),
Source: "core",
+ TraceID: "4bf92f3577b34da6a3ce929d0e0e4736",
CreatedAt: createdAt,
}
}
diff --git a/internal/tracing/audit.go b/internal/tracing/audit.go
deleted file mode 100644
index 9a3fc423c..000000000
--- a/internal/tracing/audit.go
+++ /dev/null
@@ -1,34 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you 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 tracing
-
-import "context"
-
-func PopulateAuditIDs(ctx context.Context, traceID, spanID, parentSpanID *string) {
- if traceID == nil || spanID == nil || parentSpanID == nil {
- return
- }
- if *traceID == "" && *spanID == "" {
- *traceID, *spanID = IDsFromContext(ctx)
- }
- if *parentSpanID == "" {
- if p := ParentSpanIDFromContext(ctx); p != "" {
- *parentSpanID = p
- }
- }
-}
diff --git a/internal/tracing/middleware.go b/internal/tracing/middleware.go
index 0cd183394..32a8b0587 100644
--- a/internal/tracing/middleware.go
+++ b/internal/tracing/middleware.go
@@ -23,11 +23,12 @@ import (
"go.opentelemetry.io/otel/attribute"
"go.opentelemetry.io/otel/codes"
+ "go.opentelemetry.io/otel/propagation"
"github.com/apache/airavata-custos/internal/httputil"
)
-// Middleware opens a root span per request and writes X-Trace-Id on the response.
+// Middleware opens a span per request, continuing an upstream traceparent if sent, and writes X-Trace-Id on the response.
func Middleware(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
route := r.Pattern
@@ -36,7 +37,8 @@ func Middleware(next http.Handler) http.Handler {
}
name := "http." + r.Method + " " + route
- ctx, span := Start(r.Context(), name)
+ ctx := propagation.TraceContext{}.Extract(r.Context(), propagation.HeaderCarrier(r.Header))
+ ctx, span := Start(ctx, name)
defer span.End()
span.SetAttributes(
diff --git a/internal/tracing/middleware_test.go b/internal/tracing/middleware_test.go
index 3529f45e1..0a1d54ff8 100644
--- a/internal/tracing/middleware_test.go
+++ b/internal/tracing/middleware_test.go
@@ -26,7 +26,6 @@ import (
otelcodes "go.opentelemetry.io/otel/codes"
sdktrace "go.opentelemetry.io/otel/sdk/trace"
"go.opentelemetry.io/otel/sdk/trace/tracetest"
- "go.opentelemetry.io/otel/trace/noop"
)
func setupRecordingTracer(t *testing.T) *tracetest.SpanRecorder {
@@ -39,13 +38,6 @@ func setupRecordingTracer(t *testing.T) *tracetest.SpanRecorder {
return sr
}
-func setupNoopTracer(t *testing.T) {
- t.Helper()
- prev := otel.GetTracerProvider()
- otel.SetTracerProvider(noop.NewTracerProvider())
- t.Cleanup(func() { otel.SetTracerProvider(prev) })
-}
-
func TestMiddlewareProductionEmitsRootSpanAndHeader(t *testing.T) {
sr := setupRecordingTracer(t)
@@ -110,32 +102,6 @@ func TestMiddlewareProductionEmitsRootSpanAndHeader(t *testing.T) {
}
}
-func TestMiddlewareNoopModeSetsNoHeader(t *testing.T) {
- setupNoopTracer(t)
-
- mux := http.NewServeMux()
- var (
- innerTrace string
- innerSpan string
- )
- mux.HandleFunc("GET /probe", func(w http.ResponseWriter, r *http.Request) {
- innerTrace, innerSpan = IDsFromContext(r.Context())
- w.WriteHeader(http.StatusOK)
- })
-
- handler := Middleware(mux)
- req := httptest.NewRequest(http.MethodGet, "/probe", nil)
- rec := httptest.NewRecorder()
- handler.ServeHTTP(rec, req)
-
- if hdr := rec.Header().Get("X-Trace-Id"); hdr != "" {
- t.Fatalf("expected no X-Trace-Id header in noop mode, got %q", hdr)
- }
- if innerTrace != "" || innerSpan != "" {
- t.Fatalf("expected empty IDs from noop ctx, got trace=%q span=%q", innerTrace, innerSpan)
- }
-}
-
func TestMiddleware5xxMarksSpanError(t *testing.T) {
sr := setupRecordingTracer(t)
diff --git a/internal/tracing/slog_handler.go b/internal/tracing/slog_handler.go
index 12a3cf46f..c61e1b807 100644
--- a/internal/tracing/slog_handler.go
+++ b/internal/tracing/slog_handler.go
@@ -20,6 +20,8 @@ package tracing
import (
"context"
"log/slog"
+
+ "go.opentelemetry.io/otel/trace"
)
type slogHandler struct {
@@ -35,12 +37,8 @@ func (h *slogHandler) Enabled(ctx context.Context, level slog.Level) bool {
}
func (h *slogHandler) Handle(ctx context.Context, record slog.Record) error {
- traceID, spanID := IDsFromContext(ctx)
- if traceID != "" {
- record.AddAttrs(slog.String("trace_id", traceID))
- }
- if spanID != "" {
- record.AddAttrs(slog.String("span_id", spanID))
+ if sc := trace.SpanContextFromContext(ctx); sc.IsValid() {
+ record.AddAttrs(slog.String("trace_id", sc.TraceID().String()), slog.String("span_id", sc.SpanID().String()))
}
return h.inner.Handle(ctx, record)
}
diff --git a/internal/tracing/tracing.go b/internal/tracing/tracing.go
index 92bcdebc9..940f43f45 100644
--- a/internal/tracing/tracing.go
+++ b/internal/tracing/tracing.go
@@ -31,31 +31,20 @@ import (
sdktrace "go.opentelemetry.io/otel/sdk/trace"
semconv "go.opentelemetry.io/otel/semconv/v1.27.0"
"go.opentelemetry.io/otel/trace"
- "go.opentelemetry.io/otel/trace/noop"
-)
-
-type Mode int
-
-const (
- ModeProduction Mode = iota
- ModeNoop
)
const tracerName = "custos"
type InitConfig struct {
- Mode Mode
- Logger *slog.Logger
ServiceName string
}
-func Init(cfg InitConfig) (func(context.Context) error, error) {
- if cfg.Mode == ModeNoop {
- otel.SetTracerProvider(noop.NewTracerProvider())
- otel.SetTextMapPropagator(propagation.TraceContext{})
- return func(context.Context) error { return nil }, nil
- }
+// Spans carry the audit trace ids, so they must be valid even before Init runs.
+func init() {
+ otel.SetTracerProvider(sdktrace.NewTracerProvider())
+}
+func Init(cfg InitConfig) (func(context.Context) error, error) {
serviceName := cfg.ServiceName
if serviceName == "" {
serviceName = "custos"
@@ -112,11 +101,12 @@ func Start(ctx context.Context, name string, opts ...trace.SpanStartOption) (con
return newCtx, span
}
-func ParentSpanIDFromContext(ctx context.Context) string {
+func ParentSpanIDFromContext(ctx context.Context) *string {
if p, ok := ctx.Value(parentSpanIDKey).(trace.SpanID); ok && p.IsValid() {
- return p.String()
+ s := p.String()
+ return &s
}
- return ""
+ return nil
}
func FromContext(ctx context.Context) trace.Span {
@@ -125,12 +115,8 @@ func FromContext(ctx context.Context) trace.Span {
// ContextWithTraceID attaches a saved trace id to ctx, so the spans and audit
// rows written from here on belong to the trace of the request that started it.
-// An empty traceID is normal and leaves ctx unchanged, so the caller runs with
-// no trace. An invalid one does the same but is logged since it should not happen.
+// An invalid one leaves ctx unchanged and is logged since it should not happen.
func ContextWithTraceID(ctx context.Context, traceID string) context.Context {
- if traceID == "" {
- return ctx
- }
tid, err := trace.TraceIDFromHex(traceID)
if err != nil {
slog.Warn("ignoring invalid trace id", "trace_id", traceID, "error", err)
@@ -151,7 +137,8 @@ func ContextWithTraceID(ctx context.Context, traceID string) context.Context {
func IDsFromContext(ctx context.Context) (traceID, spanID string) {
sc := trace.SpanContextFromContext(ctx)
if !sc.IsValid() {
- return "", ""
+ _, span := otel.Tracer(tracerName).Start(ctx, "detached")
+ sc = span.SpanContext()
}
return sc.TraceID().String(), sc.SpanID().String()
}
diff --git a/internal/tracing/tracing_test.go b/internal/tracing/tracing_test.go
index 71ade28db..5b991f681 100644
--- a/internal/tracing/tracing_test.go
+++ b/internal/tracing/tracing_test.go
@@ -23,33 +23,8 @@ import (
"time"
)
-func TestInitNoopProducesNonRecordingSpans(t *testing.T) {
- shutdown, err := Init(InitConfig{Mode: ModeNoop})
- if err != nil {
- t.Fatalf("Init(noop) returned error: %v", err)
- }
-
- ctx, span := Start(context.Background(), "test.noop")
- if span.IsRecording() {
- t.Fatalf("expected noop span to be non-recording")
- }
-
- traceID, spanID := IDsFromContext(ctx)
- if traceID != "" || spanID != "" {
- t.Fatalf("expected empty IDs from noop ctx, got trace=%q span=%q", traceID, spanID)
- }
-
- span.End()
-
- sctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
- defer cancel()
- if err := shutdown(sctx); err != nil {
- t.Fatalf("shutdown returned error: %v", err)
- }
-}
-
func TestInitProductionMintsValidIDs(t *testing.T) {
- shutdown, err := Init(InitConfig{Mode: ModeProduction, ServiceName: "custos-test"})
+ shutdown, err := Init(InitConfig{ServiceName: "custos-test"})
if err != nil {
t.Fatalf("Init(production) returned error: %v", err)
}
@@ -69,7 +44,7 @@ func TestInitProductionMintsValidIDs(t *testing.T) {
}
func TestStartCapturesParentSpanID(t *testing.T) {
- shutdown, err := Init(InitConfig{Mode: ModeProduction, ServiceName: "custos-test"})
+ shutdown, err := Init(InitConfig{ServiceName: "custos-test"})
if err != nil {
t.Fatalf("Init returned error: %v", err)
}
@@ -86,7 +61,7 @@ func TestStartCapturesParentSpanID(t *testing.T) {
ctx2, span2 := Start(ctx1, "inner")
defer span2.End()
- got := ParentSpanIDFromContext(ctx2)
+ got := *ParentSpanIDFromContext(ctx2)
if len(got) != 16 {
t.Fatalf("expected 16 hex chars, got %d", len(got))
}
@@ -94,7 +69,7 @@ func TestStartCapturesParentSpanID(t *testing.T) {
t.Errorf("parent mismatch: want %s got %s", span1ID.String(), got)
}
- if root := ParentSpanIDFromContext(ctx1); root != "" {
- t.Errorf("expected empty parent on root span ctx, got %s", root)
+ if root := ParentSpanIDFromContext(ctx1); root != nil {
+ t.Errorf("expected empty parent on root span ctx, got %s", *root)
}
}
diff --git a/pkg/events/integration_common_test.go b/pkg/events/integration_common_test.go
index cd836814f..33882f728 100644
--- a/pkg/events/integration_common_test.go
+++ b/pkg/events/integration_common_test.go
@@ -29,7 +29,6 @@ import (
"github.com/jmoiron/sqlx"
"github.com/apache/airavata-custos/internal/db"
- "github.com/apache/airavata-custos/internal/tracing"
)
var (
@@ -62,11 +61,6 @@ func setupTestDB(t *testing.T) *sqlx.DB {
sharedDBErr = err
return
}
- // Set up a tracer provider to embed span ids
- if _, err := tracing.Init(tracing.InitConfig{Mode: tracing.ModeProduction, ServiceName: "custos"}); err != nil {
- sharedDBErr = err
- return
- }
sharedDB = database
})
if sharedDBErr != nil {
diff --git a/pkg/events/worker.go b/pkg/events/worker.go
index 8a478a47d..b814502aa 100644
--- a/pkg/events/worker.go
+++ b/pkg/events/worker.go
@@ -264,18 +264,14 @@ func retryAudit(ctx context.Context, row models.PendingDelivery, actorID string,
"previous_attempts": row.Attempts,
"last_error": row.LastError,
})
- traceID, spanID := tracing.IDsFromContext(ctx)
return &models.AuditEvent{
- ID: uuid.NewString(),
- EventType: auditDeliveryRetried,
- EventTime: now,
- EntityID: row.ID,
- EntityType: "event_delivery",
- Details: string(body),
- Source: row.Subscriber,
- TraceID: traceID,
- SpanID: spanID,
- ParentSpanID: tracing.ParentSpanIDFromContext(ctx),
+ ID: uuid.NewString(),
+ EventType: auditDeliveryRetried,
+ EventTime: now,
+ EntityID: row.ID,
+ EntityType: "event_delivery",
+ Details: string(body),
+ Source: row.Subscriber,
}
}
diff --git a/pkg/models/audit.go b/pkg/models/audit.go
index d36adf1f0..ca92d8fce 100644
--- a/pkg/models/audit.go
+++ b/pkg/models/audit.go
@@ -29,5 +29,5 @@ type AuditEvent struct {
Source string `json:"source" db:"source"` // Subsystem that produced the event (e.g., "amie", "comanage", "slurm", "core", etc.).
TraceID string `json:"trace_id" db:"trace_id"`
SpanID string `json:"span_id" db:"span_id"`
- ParentSpanID string `json:"parent_span_id,omitempty" db:"parent_span_id"`
+ ParentSpanID *string `json:"parent_span_id,omitempty" db:"parent_span_id"`
}
diff --git a/pkg/models/trace.go b/pkg/models/trace.go
index d2ee5b353..e724a59bb 100644
--- a/pkg/models/trace.go
+++ b/pkg/models/trace.go
@@ -34,7 +34,7 @@ type TraceSummary struct {
type TraceEvent struct {
SpanID string `db:"span_id" json:"span_id"`
- ParentSpanID string `db:"parent_span_id" json:"parent_span_id,omitempty"`
+ ParentSpanID *string `db:"parent_span_id" json:"parent_span_id,omitempty"`
Source string `db:"source" json:"source"`
EventType string `db:"event_type" json:"event_type"`
EntityType string `db:"entity_type" json:"entity_type,omitempty"`
diff --git a/pkg/service/audit_event.go b/pkg/service/audit_event.go
index 719958a77..3ba3db66b 100644
--- a/pkg/service/audit_event.go
+++ b/pkg/service/audit_event.go
@@ -21,10 +21,8 @@ import (
"context"
"database/sql"
"fmt"
- "log/slog"
"github.com/apache/airavata-custos/internal/audit"
- "github.com/apache/airavata-custos/internal/tracing"
"github.com/apache/airavata-custos/pkg/models"
)
@@ -53,13 +51,6 @@ func (s *Service) CreateAuditEvent(ctx context.Context, e *models.AuditEvent) (*
e.EventTime = nowUTC()
}
- tracing.PopulateAuditIDs(ctx, &e.TraceID, &e.SpanID, &e.ParentSpanID)
- if e.TraceID == "" {
- slog.WarnContext(ctx, "audit write outside an active span",
- "event_type", e.EventType,
- "entity_id", e.EntityID,
- )
- }
if e.Source == "" {
if src := audit.SourceFromContext(ctx); src != "" {
e.Source = src
diff --git a/pkg/service/audit_event_test.go b/pkg/service/audit_event_test.go
deleted file mode 100644
index feccf9819..000000000
--- a/pkg/service/audit_event_test.go
+++ /dev/null
@@ -1,88 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you 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 service
-
-import (
- "context"
- "testing"
-
- "go.opentelemetry.io/otel"
- sdktrace "go.opentelemetry.io/otel/sdk/trace"
-
- "github.com/apache/airavata-custos/internal/tracing"
- "github.com/apache/airavata-custos/pkg/models"
-)
-
-func TestAuditEventTraceIDs_PopulatedFromActiveSpan(t *testing.T) {
- prev := otel.GetTracerProvider()
- tp := sdktrace.NewTracerProvider()
- otel.SetTracerProvider(tp)
- t.Cleanup(func() { otel.SetTracerProvider(prev) })
-
- ctx, span := tracing.Start(context.Background(), "test.root")
- defer span.End()
-
- wantTrace := span.SpanContext().TraceID()
- wantSpan := span.SpanContext().SpanID()
-
- e := &models.AuditEvent{EventType: "X", EntityID: "y"}
- tracing.PopulateAuditIDs(ctx, &e.TraceID, &e.SpanID, &e.ParentSpanID)
-
- if e.TraceID != wantTrace.String() {
- t.Fatalf("trace_id mismatch: got %s want %s", e.TraceID, wantTrace.String())
- }
- if e.SpanID != wantSpan.String() {
- t.Fatalf("span_id mismatch: got %s want %s", e.SpanID, wantSpan.String())
- }
-}
-
-func TestAuditEventTraceIDs_NilWhenNoSpan(t *testing.T) {
- e := &models.AuditEvent{EventType: "X", EntityID: "y"}
- tracing.PopulateAuditIDs(context.Background(), &e.TraceID, &e.SpanID, &e.ParentSpanID)
-
- if e.TraceID != "" || e.SpanID != "" {
- t.Fatalf("expected empty trace/span IDs, got trace=%s span=%s", e.TraceID, e.SpanID)
- }
-}
-
-func TestAuditEventTraceIDs_NotOverwrittenWhenPreset(t *testing.T) {
- prev := otel.GetTracerProvider()
- tp := sdktrace.NewTracerProvider()
- otel.SetTracerProvider(tp)
- t.Cleanup(func() { otel.SetTracerProvider(prev) })
-
- ctx, span := tracing.Start(context.Background(), "test.root")
- defer span.End()
-
- preset := "0102030405060708090a0b0c0d0e0f10"
- presetSpan := "0102030405060708"
- e := &models.AuditEvent{
- EventType: "X",
- EntityID: "y",
- TraceID: preset,
- SpanID: presetSpan,
- }
- tracing.PopulateAuditIDs(ctx, &e.TraceID, &e.SpanID, &e.ParentSpanID)
-
- if e.TraceID != preset {
- t.Fatalf("preset trace_id was overwritten: got %s want %s", e.TraceID, preset)
- }
- if e.SpanID != presetSpan {
- t.Fatalf("preset span_id was overwritten: got %s want %s", e.SpanID, presetSpan)
- }
-}