You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airavata.apache.org by la...@apache.org on 2026/10/09 20:03:01 UTC
(airavata-custos) branch master updated: Make trace ids mandatory and report trace status (#607)
This is an automated email from the ASF dual-hosted git repository.
lahirujayathilake pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/airavata-custos.git
The following commit(s) were added to refs/heads/master by this push:
new b9c205ee4 Make trace ids mandatory and report trace status (#607)
b9c205ee4 is described below
commit b9c205ee4f7f2106cb8f1b820c7b0b8f85f1311d
Author: Yasith Jayawardana <ya...@icloud.com>
AuthorDate: Fri Oct 9 16:02:55 2026 -0400
Make trace ids mandatory and report trace status (#607)
* Make trace ids mandatory and report trace status
* Warn once when a SLURM allocation has no resources
---
api/core.openapi.yaml | 5 +
cmd/server/bootstrap.go | 5 +-
cmd/server/main.go | 7 +-
.../ACCESS/AMIE-Processor/amieclient/client.go | 14 +-
.../db/migrations/000001_initial_schema.up.sql | 2 +
.../ACCESS/AMIE-Processor/handler/held_reply.go | 3 +
.../handler/held_reply_integration_test.go | 11 +-
.../AMIE-Processor/handler/integration_common.go | 3 +-
.../request_person_merge_integration_test.go | 4 +-
.../handler/request_project_create.go | 2 +-
connectors/ACCESS/AMIE-Processor/model/event.go | 2 +
.../AMIE-Processor/pipeline/integration_common.go | 34 +----
.../ACCESS/AMIE-Processor/pkg/amie/loader.go | 2 +-
.../ACCESS/AMIE-Processor/service/audit_service.go | 11 --
.../AMIE-Processor/service/audit_service_test.go | 57 --------
.../ACCESS/AMIE-Processor/store/event_store.go | 8 +-
.../AMIE-Processor/store/packet_audit_store.go | 2 +-
connectors/ACCESS/AMIE-Processor/worker/poller.go | 48 ++++---
.../worker/poller_integration_test.go | 14 ++
.../ACCESS/AMIE-Processor/worker/processor.go | 13 +-
.../worker/processor_integration_test.go | 13 +-
.../internal/operations/ensure_posix_account.go | 14 +-
.../internal/operations/orchestrator_test.go | 6 +-
.../internal/subscribers/allocation_group.go | 11 ++
.../internal/subscribers/cluster_user.go | 14 +-
.../Identity-Provisioner/pkg/comanage/loader.go | 2 +-
.../internal/notify/notifier_integration_test.go | 3 +-
.../internal/subscribers/account.go | 36 ++---
.../internal/subscribers/association_writer.go | 17 ++-
.../internal/subscribers/deprovision.go | 8 +-
.../internal/subscribers/members.go | 26 ++--
.../internal/subscribers/reconciler.go | 30 +++--
.../internal/subscribers/subscriber_trace_test.go | 6 +-
.../Usage-Monitor/internal/smonitor/smonitor.go | 35 ++---
.../internal/subscribers/subscriber.go | 5 +-
internal/db/migrations/000005_audit.up.sql | 6 +-
internal/db/migrations/000007_event_bus.up.sql | 3 +-
internal/server/audit_handlers.go | 3 +-
internal/server/audit_handlers_test.go | 6 +-
.../server/cluster_account_integration_test.go | 3 +-
internal/server/event_delivery_integration_test.go | 5 +-
internal/server/privilege_integration_test.go | 14 +-
internal/server/server.go | 18 ++-
internal/server/server_trace_test.go | 35 +++--
internal/server/testidentity_test.go | 4 +-
internal/store/audit_event_store.go | 8 ++
internal/store/audit_trace_store.go | 149 +++++++++++----------
.../store/audit_trace_store_integration_test.go | 96 +++++++++++++
internal/store/audit_trace_store_test.go | 31 ++---
internal/store/event_delivery_store.go | 8 +-
.../store/event_delivery_store_integration_test.go | 2 +
internal/tracing/audit.go | 17 ++-
internal/tracing/middleware.go | 33 +++--
internal/tracing/middleware_test.go | 64 +++------
internal/tracing/slog_handler.go | 10 +-
internal/tracing/status.go | 78 +++++------
internal/tracing/status_test.go | 67 +++------
internal/tracing/tracing.go | 83 +++++-------
internal/tracing/tracing_test.go | 37 +----
.../{audit.go => tracingtest/tracingtest.go} | 25 ++--
pkg/events/bus.go | 12 +-
pkg/events/bus_test.go | 27 ++--
pkg/events/integration_common_test.go | 6 -
pkg/events/worker.go | 12 +-
pkg/identity/identity.go | 4 +
pkg/identity/middleware.go | 4 +-
pkg/models/audit.go | 2 +-
pkg/models/event_delivery.go | 1 +
pkg/models/trace.go | 2 +-
pkg/service/audit_event.go | 22 ---
pkg/service/audit_event_test.go | 22 ---
pkg/service/identity_resolver.go | 6 +-
pkg/service/integration_common_test.go | 3 +-
pkg/service/user_onboard.go | 2 +-
pkg/service/user_role.go | 8 +-
75 files changed, 693 insertions(+), 718 deletions(-)
diff --git a/api/core.openapi.yaml b/api/core.openapi.yaml
index 54dfd5baa..e6eb90779 100644
--- a/api/core.openapi.yaml
+++ b/api/core.openapi.yaml
@@ -485,6 +485,9 @@ definitions:
source:
description: Subsystem that published the event.
type: string
+ span_id:
+ description: Publisher's span, the parent of its deliveries.
+ type: string
trace_id:
type: string
type: object
@@ -1147,6 +1150,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..70a056e22 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,7 +38,9 @@ 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)
+ slog.WarnContext(ctx, "bootstrap super_admin failed", "email", email, "error", err)
}
}
diff --git a/cmd/server/main.go b/cmd/server/main.go
index 375faf5a8..a5619919b 100644
--- a/cmd/server/main.go
+++ b/cmd/server/main.go
@@ -102,12 +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",
})
@@ -161,7 +156,7 @@ func run() error {
// identity.Middleware sits in front of the router-backed server, so every
// gated route sees a verified caller + privilege set on ctx.
authed := identity.Middleware(verifier, svc, router.PublicPaths(), srv)
- handler := server.LoggingMiddleware(authed)
+ handler := server.LoggingMiddleware(tracing.Middleware(authed))
httpServer := &http.Server{
Addr: addr,
diff --git a/connectors/ACCESS/AMIE-Processor/amieclient/client.go b/connectors/ACCESS/AMIE-Processor/amieclient/client.go
index bd8734a3e..347d7b667 100644
--- a/connectors/ACCESS/AMIE-Processor/amieclient/client.go
+++ b/connectors/ACCESS/AMIE-Processor/amieclient/client.go
@@ -70,7 +70,7 @@ func (c *Client) FetchInProgressPackets(ctx context.Context) ([]map[string]any,
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
if err != nil {
- slog.Warn("failed to create AMIE fetch request", "error", err)
+ slog.WarnContext(ctx, "failed to create AMIE fetch request", "error", err)
return []map[string]any{}, nil
}
req.Header.Set("XA-SITE", c.siteCode)
@@ -79,23 +79,23 @@ func (c *Client) FetchInProgressPackets(ctx context.Context) ([]map[string]any,
resp, err := c.httpClient.Do(req)
if err != nil {
- slog.Warn("AMIE fetch request failed", "error", err)
+ slog.WarnContext(ctx, "AMIE fetch request failed", "error", err)
return []map[string]any{}, nil
}
defer resp.Body.Close()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
- slog.Warn("AMIE fetch returned non-2xx status", "status", resp.StatusCode)
+ slog.WarnContext(ctx, "AMIE fetch returned non-2xx status", "status", resp.StatusCode)
return []map[string]any{}, nil
}
body, err := io.ReadAll(resp.Body)
if err != nil {
- slog.Warn("failed to read AMIE response body", "error", err)
+ slog.WarnContext(ctx, "failed to read AMIE response body", "error", err)
return []map[string]any{}, nil
}
- return parsePacketsFromResponse(body)
+ return parsePacketsFromResponse(ctx, body)
}
// parsePacketsFromResponse extracts a slice of packet maps from the AMIE API
@@ -104,7 +104,7 @@ func (c *Client) FetchInProgressPackets(ctx context.Context) ([]map[string]any,
// 1. A root JSON array of packet objects.
// 2. A root JSON object with a "result" key containing an array.
// 3. A root JSON object (single packet or envelope with scalar "result").
-func parsePacketsFromResponse(body []byte) ([]map[string]any, error) {
+func parsePacketsFromResponse(ctx context.Context, body []byte) ([]map[string]any, error) {
// Try as a JSON array first.
var arr []map[string]any
if err := json.Unmarshal(body, &arr); err == nil {
@@ -114,7 +114,7 @@ func parsePacketsFromResponse(body []byte) ([]map[string]any, error) {
// Try as a JSON object.
var obj map[string]any
if err := json.Unmarshal(body, &obj); err != nil {
- slog.Error("failed to parse AMIE response JSON", "error", err)
+ slog.ErrorContext(ctx, "failed to parse AMIE response JSON", "error", err)
return []map[string]any{}, nil
}
diff --git a/connectors/ACCESS/AMIE-Processor/db/migrations/000001_initial_schema.up.sql b/connectors/ACCESS/AMIE-Processor/db/migrations/000001_initial_schema.up.sql
index a469f9e91..e85f10545 100644
--- a/connectors/ACCESS/AMIE-Processor/db/migrations/000001_initial_schema.up.sql
+++ b/connectors/ACCESS/AMIE-Processor/db/migrations/000001_initial_schema.up.sql
@@ -53,6 +53,8 @@ CREATE TABLE IF NOT EXISTS amie_processing_events
finished_at TIMESTAMPTZ(6) NULL,
last_error TEXT NULL,
next_retry_at TIMESTAMPTZ(6) NULL DEFAULT NULL,
+ trace_id CHAR(32) NOT NULL CHECK (trace_id <> ''),
+ span_id CHAR(16) NOT NULL CHECK (span_id <> ''),
PRIMARY KEY (id),
CONSTRAINT fk_amie_events_packet FOREIGN KEY (packet_id) REFERENCES amie_packets (id) ON DELETE CASCADE,
CONSTRAINT uq_amie_events_packet_type UNIQUE (packet_id, type)
diff --git a/connectors/ACCESS/AMIE-Processor/handler/held_reply.go b/connectors/ACCESS/AMIE-Processor/handler/held_reply.go
index 9df1efb6f..1bb15b924 100644
--- a/connectors/ACCESS/AMIE-Processor/handler/held_reply.go
+++ b/connectors/ACCESS/AMIE-Processor/handler/held_reply.go
@@ -28,6 +28,7 @@ import (
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/connectors/ACCESS/AMIE-Processor/store"
+ "github.com/apache/airavata-custos/internal/tracing"
"github.com/apache/airavata-custos/pkg/events"
"github.com/apache/airavata-custos/pkg/models"
)
@@ -93,6 +94,8 @@ func (h *HeldReplies) denied(ctx context.Context, cu models.ComputeClusterUser)
// Each packet gets its own transaction, so a failed send leaves only that
// packet waiting for the bus to retry.
func (h *HeldReplies) answer(ctx context.Context, clusterUserID string, status model.PacketStatus, failure map[string]any) error {
+ ctx, span := tracing.Start(ctx, "amie.held_reply")
+ defer span.End()
packets, err := h.packets.ListWaitingFor(ctx, clusterUserID)
if err != nil {
return fmt.Errorf("list held packets: %w", err)
diff --git a/connectors/ACCESS/AMIE-Processor/handler/held_reply_integration_test.go b/connectors/ACCESS/AMIE-Processor/handler/held_reply_integration_test.go
index 4cada00b1..62c47f10d 100644
--- a/connectors/ACCESS/AMIE-Processor/handler/held_reply_integration_test.go
+++ b/connectors/ACCESS/AMIE-Processor/handler/held_reply_integration_test.go
@@ -30,6 +30,8 @@ import (
"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/internal/tracing/tracingtest"
"github.com/apache/airavata-custos/pkg/models"
coreservice "github.com/apache/airavata-custos/pkg/service"
)
@@ -81,11 +83,16 @@ func TestHeldReplies_ApprovalSendsHeldReplyOnce(t *testing.T) {
amie := &fakeAmieClient{}
held := NewHeldReplies(database, store.NewPacketStore(database), amie, newTestAuditService(database))
account := reviewedAccount(t, svc, *pkt.HeldFor)
+ ctx, delivery := tracing.Start(context.Background(), "delivery")
for range 2 {
- if err := held.approved(context.Background(), account); err != nil {
+ if err := held.approved(ctx, account); err != nil {
t.Fatalf("approved: %v", err)
}
}
+ var chained int
+ if err := database.Get(&chained, "SELECT COUNT(*) FROM audit_events WHERE event_type = 'REPLY_SENT' AND parent_span_id = $1", delivery.SpanContext().SpanID().String()); err != nil || chained != 1 {
+ t.Errorf("REPLY_SENT rows under the delivery span: %d (%v), want 1", chained, err)
+ }
if len(amie.Replies) != 1 || amie.lastReplyType() != "notify_account_create" {
t.Fatalf("expected one notify_account_create, got %+v", amie.Replies)
@@ -142,7 +149,7 @@ func TestHeldReplies_DenialSendsFailureWithReason(t *testing.T) {
if err := database.Get(&reviewer, "SELECT id FROM users LIMIT 1"); err != nil {
t.Fatalf("pick a reviewer: %v", err)
}
- if _, err := svc.DenyComputeClusterUser(context.Background(), *pkt.HeldFor, reviewer, "not on the collaborator list"); err != nil {
+ if _, err := svc.DenyComputeClusterUser(tracingtest.Context(), *pkt.HeldFor, reviewer, "not on the collaborator list"); err != nil {
t.Fatalf("deny: %v", err)
}
diff --git a/connectors/ACCESS/AMIE-Processor/handler/integration_common.go b/connectors/ACCESS/AMIE-Processor/handler/integration_common.go
index a6e9028cb..61c5a657a 100644
--- a/connectors/ACCESS/AMIE-Processor/handler/integration_common.go
+++ b/connectors/ACCESS/AMIE-Processor/handler/integration_common.go
@@ -38,6 +38,7 @@ import (
"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/internal/tracing/tracingtest"
"github.com/apache/airavata-custos/pkg/events"
coreservice "github.com/apache/airavata-custos/pkg/service"
)
@@ -233,7 +234,7 @@ func approvePendingAccounts(t *testing.T, database *sqlx.DB, svc *coreservice.Se
t.Fatalf("pick a reviewer: %v", err)
}
for _, id := range ids {
- if _, err := svc.ApproveComputeClusterUser(context.Background(), id, reviewer); err != nil {
+ if _, err := svc.ApproveComputeClusterUser(tracingtest.Context(), id, reviewer); err != nil {
t.Fatalf("approve %s: %v", id, err)
}
}
diff --git a/connectors/ACCESS/AMIE-Processor/handler/request_person_merge_integration_test.go b/connectors/ACCESS/AMIE-Processor/handler/request_person_merge_integration_test.go
index f991c99b7..cf7e0faba 100644
--- a/connectors/ACCESS/AMIE-Processor/handler/request_person_merge_integration_test.go
+++ b/connectors/ACCESS/AMIE-Processor/handler/request_person_merge_integration_test.go
@@ -566,8 +566,8 @@ func TestRequestPersonMerge_ReprocessesPacketsHeldOnRetiringAccount(t *testing.T
t.Fatalf("hold packet: %v", err)
}
eventID := uuid.NewString()
- if _, err := database.Exec("INSERT INTO amie_processing_events (id, packet_id, type, status) VALUES ($1, $2, $3, $4)",
- eventID, held.ID, held.Type, model.ProcessingStatusSucceeded); err != nil {
+ if _, err := database.Exec("INSERT INTO amie_processing_events (id, packet_id, type, status, trace_id, span_id) VALUES ($1, $2, $3, $4, $5, $6)",
+ eventID, held.ID, held.Type, model.ProcessingStatusSucceeded, "4bf92f3577b34da6a3ce929d0e0e4736", "00f067aa0ba902b7"); err != nil {
t.Fatalf("insert processing event: %v", err)
}
diff --git a/connectors/ACCESS/AMIE-Processor/handler/request_project_create.go b/connectors/ACCESS/AMIE-Processor/handler/request_project_create.go
index a74375c59..701880d4c 100644
--- a/connectors/ACCESS/AMIE-Processor/handler/request_project_create.go
+++ b/connectors/ACCESS/AMIE-Processor/handler/request_project_create.go
@@ -257,7 +257,7 @@ func (h *RequestProjectCreateHandler) ensureResourceMappings(ctx context.Context
for _, name := range getResourceList(body) {
resource, err := h.svc.GetComputeAllocationResourceByNameAndCluster(ctx, name, h.clusterID)
if errors.Is(err, service.ErrNotFound) {
- slog.Warn("amie: packet resource not in the cluster catalog, mapping skipped",
+ slog.WarnContext(ctx, "amie: packet resource not in the cluster catalog, mapping skipped",
"resource", name, "allocation_id", allocation.ID)
continue
}
diff --git a/connectors/ACCESS/AMIE-Processor/model/event.go b/connectors/ACCESS/AMIE-Processor/model/event.go
index 3dc050a8d..1ff74ad18 100644
--- a/connectors/ACCESS/AMIE-Processor/model/event.go
+++ b/connectors/ACCESS/AMIE-Processor/model/event.go
@@ -50,6 +50,8 @@ type ProcessingEvent struct {
FinishedAt *time.Time `db:"finished_at" json:"finished_at,omitempty"`
LastError *string `db:"last_error" json:"last_error,omitempty"`
NextRetryAt *time.Time `db:"next_retry_at" json:"next_retry_at,omitempty"`
+ TraceID string `db:"trace_id" json:"trace_id"` // Poll pass that ingested the packet.
+ SpanID string `db:"span_id" json:"span_id"`
}
// EventWithPacket is a read-only projection that joins a processing event with
diff --git a/connectors/ACCESS/AMIE-Processor/pipeline/integration_common.go b/connectors/ACCESS/AMIE-Processor/pipeline/integration_common.go
index 84699d604..c3d7bc866 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"
@@ -32,8 +31,6 @@ import (
"github.com/jmoiron/sqlx"
"github.com/prometheus/client_golang/prometheus"
- "go.opentelemetry.io/otel"
- sdktrace "go.opentelemetry.io/otel/sdk/trace"
"github.com/apache/airavata-custos/connectors/ACCESS/AMIE-Processor/amieclient"
amiedb "github.com/apache/airavata-custos/connectors/ACCESS/AMIE-Processor/db"
@@ -52,11 +49,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 +97,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
@@ -279,9 +264,11 @@ func (p *testPipeline) startApproving() {
len(ids) == 0 || p.db.GetContext(p.ctx, &reviewer, "SELECT id FROM users LIMIT 1") != nil {
continue
}
+ ctx, span := tracing.Start(p.ctx, "test.approve")
for _, id := range ids {
- _, _ = p.svc.ApproveComputeClusterUser(p.ctx, id, reviewer)
+ _, _ = p.svc.ApproveComputeClusterUser(ctx, id, reviewer)
}
+ span.End()
}
}()
}
@@ -289,15 +276,6 @@ func (p *testPipeline) startApproving() {
func (p *testPipeline) stop() {
p.cancel()
p.wg.Wait()
- flushTracing()
-}
-
-func flushTracing() {
- if tp, ok := otel.GetTracerProvider().(*sdktrace.TracerProvider); ok {
- ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
- defer cancel()
- _ = tp.ForceFlush(ctx)
- }
}
func (p *testPipeline) fireScenario(t *testing.T, scenarioType string) {
diff --git a/connectors/ACCESS/AMIE-Processor/pkg/amie/loader.go b/connectors/ACCESS/AMIE-Processor/pkg/amie/loader.go
index 338310eeb..b49e25bb5 100644
--- a/connectors/ACCESS/AMIE-Processor/pkg/amie/loader.go
+++ b/connectors/ACCESS/AMIE-Processor/pkg/amie/loader.go
@@ -49,7 +49,7 @@ import (
)
func init() {
- tracing.RegisterTerminalMarkers("amie", "TRANSACTION_COMPLETE")
+ tracing.RegisterMarkers("amie", "PACKET_RECEIVED", "TRANSACTION_COMPLETE", "REPLY_SENT", "REPLY_HELD")
}
const connectorName = "amie"
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/event_store.go b/connectors/ACCESS/AMIE-Processor/store/event_store.go
index f6eb07648..84e7312c0 100644
--- a/connectors/ACCESS/AMIE-Processor/store/event_store.go
+++ b/connectors/ACCESS/AMIE-Processor/store/event_store.go
@@ -58,7 +58,7 @@ func (s *pgEventStore) FindByID(ctx context.Context, id string) (*model.Processi
}
func (s *pgEventStore) FindTop50EventsToProcess(ctx context.Context, statuses []model.ProcessingStatus, now time.Time) ([]model.EventWithPacket, error) {
- query := `SELECT e.id, e.packet_id, e.type, e.status, e.attempts, e.created_at, e.started_at, e.finished_at, e.last_error, e.next_retry_at,
+ query := `SELECT e.id, e.packet_id, e.type, e.status, e.attempts, e.created_at, e.started_at, e.finished_at, e.last_error, e.next_retry_at, e.trace_id, e.span_id,
p.amie_id AS packet_amie_id, p.type AS packet_type, p.raw_json AS packet_raw_json
FROM amie_processing_events e
JOIN amie_packets p ON e.packet_id = p.id
@@ -85,11 +85,11 @@ func (s *pgEventStore) FindTop50EventsToProcess(ctx context.Context, statuses []
func (s *pgEventStore) Save(ctx context.Context, tx *sql.Tx, e *model.ProcessingEvent) error {
_, err := tx.ExecContext(ctx,
- `INSERT INTO amie_processing_events (id, packet_id, type, status, attempts, created_at, started_at, finished_at, last_error, next_retry_at)
- VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)`,
+ `INSERT INTO amie_processing_events (id, packet_id, type, status, attempts, created_at, started_at, finished_at, last_error, next_retry_at, trace_id, span_id)
+ VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)`,
e.ID, e.PacketID, e.Type, e.Status, e.Attempts,
e.CreatedAt, e.StartedAt, e.FinishedAt,
- e.LastError, e.NextRetryAt)
+ e.LastError, e.NextRetryAt, e.TraceID, e.SpanID)
return err
}
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..4c015bcc4 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,25 +99,33 @@ func (p *Poller) Run(ctx context.Context) {
}
func (p *Poller) pollForPackets(ctx context.Context) {
- slog.Info("Polling for new AMIE packets...")
+ for _, packetNode := range p.fetchPackets(ctx) {
+ p.processIndividualPacket(ctx, packetNode)
+ }
+}
+
+// fetchPackets runs in its own span, so each packet below starts a trace of its own.
+func (p *Poller) fetchPackets(ctx context.Context) []map[string]any {
+ ctx, span := tracing.Start(ctx, "amie.poll")
+ defer span.End()
+ slog.InfoContext(ctx, "Polling for new AMIE packets...")
packets, err := p.client.FetchInProgressPackets(ctx)
if err != nil {
- slog.Error("failed to fetch AMIE packets", "error", err)
- return
+ slog.ErrorContext(ctx, "failed to fetch AMIE packets", "error", err)
+ return nil
}
if len(packets) == 0 {
- slog.Debug("no new AMIE packets found")
- return
+ slog.DebugContext(ctx, "no new AMIE packets found")
+ return nil
}
- slog.Info("fetched AMIE packets", "count", len(packets))
+ slog.InfoContext(ctx, "fetched AMIE packets", "count", len(packets))
p.metrics.RecordPollerFetch(len(packets))
-
- for _, packetNode := range packets {
- p.processIndividualPacket(ctx, packetNode)
- }
+ return packets
}
func (p *Poller) processIndividualPacket(ctx context.Context, packetNode map[string]any) {
+ ctx, span := tracing.Start(ctx, "amie.ingest")
+ defer span.End()
// Extract amie_id from header.packet_rec_id.
amiePacketRecID := int64(-1)
if header, ok := packetNode["header"].(map[string]any); ok {
@@ -135,7 +144,7 @@ func (p *Poller) processIndividualPacket(ctx context.Context, packetNode map[str
}
if amiePacketRecID < 0 || packetType == "" {
- slog.Warn("skipping packet with missing header fields", "raw", packetNode)
+ slog.WarnContext(ctx, "skipping packet with missing header fields", "raw", packetNode)
return
}
@@ -144,11 +153,11 @@ func (p *Poller) processIndividualPacket(ctx context.Context, packetNode map[str
// Deduplication check.
existing, err := p.packetStore.FindByAmieID(ctx, amiePacketRecID)
if err != nil {
- logger.Error("failed to check for existing packet", "error", err)
+ logger.ErrorContext(ctx, "failed to check for existing packet", "error", err)
return
}
if existing != nil {
- logger.Debug("packet already exists, skipping")
+ logger.DebugContext(ctx, "packet already exists, skipping")
return
}
@@ -170,24 +179,29 @@ func (p *Poller) processIndividualPacket(ctx context.Context, packetNode map[str
}
// Create processing event.
- eventID := uuid.NewString()
+ traceID, spanID, err := tracing.IDsFromContext(ctx)
+ if err != nil {
+ return err
+ }
event := &model.ProcessingEvent{
- ID: eventID,
+ ID: uuid.NewString(),
PacketID: newPacket.ID,
Type: model.EventTypeDecodePacket,
Status: model.ProcessingStatusNew,
CreatedAt: now,
+ TraceID: traceID,
+ SpanID: spanID,
}
if err := p.eventStore.Save(ctx, tx, event); err != nil {
return fmt.Errorf("save processing event: %w", err)
}
- logger.Info("persisted new AMIE packet", "packetId", newPacket.ID)
+ logger.InfoContext(ctx, "persisted new AMIE packet", "packetId", newPacket.ID)
return nil
})
if err != nil {
- logger.Error("failed to persist packet", "error", err)
+ logger.ErrorContext(ctx, "failed to persist packet", "error", err)
return
}
diff --git a/connectors/ACCESS/AMIE-Processor/worker/poller_integration_test.go b/connectors/ACCESS/AMIE-Processor/worker/poller_integration_test.go
index 049169e69..53e813b8f 100644
--- a/connectors/ACCESS/AMIE-Processor/worker/poller_integration_test.go
+++ b/connectors/ACCESS/AMIE-Processor/worker/poller_integration_test.go
@@ -208,3 +208,17 @@ func TestPoller_DisabledRunReturnsImmediately(t *testing.T) {
t.Errorf("disabled poller called client %d times, want 0", stub.callCount())
}
}
+
+// Packets fetched in one poll each start a trace of their own.
+func TestPoller_EachPacketGetsItsOwnTrace(t *testing.T) {
+ database := setupTestDB(t)
+ stub := &stubAmieClient{responses: [][]map[string]any{{makePollerPacket(1101, "request_project_create"), makePollerPacket(1102, "request_account_create")}}}
+ p, _ := newPoller(t, stub)
+
+ p.pollForPackets(context.Background())
+
+ var traces []string
+ if err := database.Select(&traces, "SELECT DISTINCT trace_id FROM amie_processing_events"); err != nil || len(traces) != 2 {
+ t.Fatalf("traces = %v (%v), want one per packet", traces, err)
+ }
+}
diff --git a/connectors/ACCESS/AMIE-Processor/worker/processor.go b/connectors/ACCESS/AMIE-Processor/worker/processor.go
index 691c3f7ea..e97b3ce11 100644
--- a/connectors/ACCESS/AMIE-Processor/worker/processor.go
+++ b/connectors/ACCESS/AMIE-Processor/worker/processor.go
@@ -138,6 +138,7 @@ func (p *Processor) processPendingEvents(ctx context.Context) {
slog.Info("processing pending events", "count", len(events))
for _, ewp := range events {
+ ctx := tracing.ContextWithSpanContext(ctx, ewp.TraceID, ewp.SpanID)
logger := slog.With(
"eventId", ewp.ID,
"packetId", ewp.PacketID,
@@ -148,9 +149,9 @@ func (p *Processor) processPendingEvents(ctx context.Context) {
stopTimer := p.metrics.StartProcessingTimer()
if err := p.executeInTransaction(ctx, ewp); err != nil {
- logger.Error("event processing failed", "error", err)
+ logger.ErrorContext(ctx, "event processing failed", "error", err)
if recordErr := p.recordFailureInNewTransaction(ctx, ewp.ID, err); recordErr != nil {
- logger.Error("failed to record processing failure", "error", recordErr)
+ logger.ErrorContext(ctx, "failed to record processing failure", "error", recordErr)
}
}
@@ -184,7 +185,7 @@ func (p *Processor) executeInTransaction(ctx context.Context, ewp model.EventWit
}()
return custosdb.TxFn(ctx, p.db, func(tx *sql.Tx) error {
- slog.Info("Processing event",
+ slog.InfoContext(ctx, "Processing event",
"eventId", ewp.ID,
"type", ewp.Type,
"attempt", ewp.Attempts+1,
@@ -253,7 +254,7 @@ func (p *Processor) recordFailureInNewTransaction(ctx context.Context, eventID s
return fmt.Errorf("find event for failure recording: %w", err)
}
if event == nil {
- slog.Warn("event not found for failure recording", "eventId", eventID)
+ slog.WarnContext(ctx, "event not found for failure recording", "eventId", eventID)
return nil
}
@@ -286,7 +287,7 @@ func (p *Processor) recordFailureInNewTransaction(ctx context.Context, eventID s
p.metrics.RecordRetry()
p.metrics.RecordPacketProcessed(packet.Type, "retry_scheduled")
- slog.Warn("event failed, scheduling retry",
+ slog.WarnContext(ctx, "event failed, scheduling retry",
"eventId", eventID,
"attempt", effectiveAttempts,
"nextRetryAt", nextRetry,
@@ -295,7 +296,7 @@ func (p *Processor) recordFailureInNewTransaction(ctx context.Context, eventID s
event.Status = model.ProcessingStatusPermanentlyFailed
event.NextRetryAt = nil
p.metrics.RecordPacketProcessed(packet.Type, "permanently_failed")
- slog.Error("event permanently failed after max attempts", "eventId", eventID)
+ slog.ErrorContext(ctx, "event permanently failed after max attempts", "eventId", eventID)
packet.Status = model.PacketStatusFailed
packet.Retries = effectiveAttempts
diff --git a/connectors/ACCESS/AMIE-Processor/worker/processor_integration_test.go b/connectors/ACCESS/AMIE-Processor/worker/processor_integration_test.go
index f61dc2674..9e7da347b 100644
--- a/connectors/ACCESS/AMIE-Processor/worker/processor_integration_test.go
+++ b/connectors/ACCESS/AMIE-Processor/worker/processor_integration_test.go
@@ -57,6 +57,9 @@ func (f *fakeRouter) Route(_ context.Context, _ *sql.Tx, _ map[string]any, _ *mo
return nil
}
+// The poll pass that seeded events are ingested under.
+const pollTraceID, pollSpanID = "4bf92f3577b34da6a3ce929d0e0e4736", "00f067aa0ba902b7"
+
// seedNewEvent inserts an amie_packets row + amie_processing_events row both
// in NEW state. Returns the IDs so tests can refetch state directly.
func seedNewEvent(t *testing.T, database *sqlx.DB) (packetID, eventID string) {
@@ -75,9 +78,9 @@ func seedNewEvent(t *testing.T, database *sqlx.DB) (packetID, eventID string) {
}
if _, err := database.Exec(
- `INSERT INTO amie_processing_events (id, packet_id, type, status, attempts, created_at)
- VALUES ($1, $2, $3, $4, $5, $6)`,
- eventID, packetID, string(model.EventTypeDecodePacket), string(model.ProcessingStatusNew), 0, now,
+ `INSERT INTO amie_processing_events (id, packet_id, type, status, attempts, created_at, trace_id, span_id)
+ VALUES ($1, $2, $3, $4, $5, $6, $7, $8)`,
+ eventID, packetID, string(model.EventTypeDecodePacket), string(model.ProcessingStatusNew), 0, now, pollTraceID, pollSpanID,
); err != nil {
t.Fatalf("seed event: %v", err)
}
@@ -163,6 +166,10 @@ func TestProcessor_SuccessFirstAttempt(t *testing.T) {
if got := countRows(t, database, "amie_processing_errors"); got != 0 {
t.Errorf("error rows: got %d, want 0", got)
}
+ var trace, parent string
+ if err := database.QueryRow("SELECT trace_id, parent_span_id FROM audit_events WHERE event_type = 'PACKET_RECEIVED'").Scan(&trace, &parent); err != nil || trace != pollTraceID || parent != pollSpanID {
+ t.Errorf("PACKET_RECEIVED trace %q parent %q (%v), want the poll pass %q %q", trace, parent, err, pollTraceID, pollSpanID)
+ }
}
// TestProcessor_FailTwiceSucceedThird, three attempts total; first two fail,
diff --git a/connectors/COmanage/Identity-Provisioner/internal/operations/ensure_posix_account.go b/connectors/COmanage/Identity-Provisioner/internal/operations/ensure_posix_account.go
index 04de86c0e..06ed91a7f 100644
--- a/connectors/COmanage/Identity-Provisioner/internal/operations/ensure_posix_account.go
+++ b/connectors/COmanage/Identity-Provisioner/internal/operations/ensure_posix_account.go
@@ -55,13 +55,13 @@ func (o *Orchestrator) ensurePOSIXAccountImpl(ctx context.Context, cu *models.Co
return err
}
log = log.With("comanage_person_id", personID)
- log.Info("comanage: CoPerson resolved", "created", created)
+ log.InfoContext(ctx, "comanage: CoPerson resolved", "created", created)
if created {
o.audit(ctx, cu, "ComanageCoPersonCreated", fmt.Sprintf("comanage_id=%s email=%s", personID, user.Email))
}
if err := o.storePersonID(ctx, cu.UserID, personID); err != nil {
- log.Warn("comanage: failed to store CoPerson id", "err", err)
+ log.WarnContext(ctx, "comanage: failed to store CoPerson id", "err", err)
}
if composite == nil {
@@ -146,10 +146,10 @@ func (o *Orchestrator) ensurePOSIXAccountImpl(ctx context.Context, cu *models.Co
o.dlq(ctx, cu, "put_composite", err)
return err
}
- log.Info("comanage: UnixClusterAccount attached", "username", cu.LocalUsername, "uid", uidInt, "co_group_id", coGroupID)
+ log.InfoContext(ctx, "comanage: UnixClusterAccount attached", "username", cu.LocalUsername, "uid", uidInt, "co_group_id", coGroupID)
o.audit(ctx, cu, "ComanageClusterAccountAttached", fmt.Sprintf("comanage_id=%s username=%s uid=%d", personID, cu.LocalUsername, uidInt))
if err := o.core.MarkComputeClusterUserProvisioned(ctx, cu.ID); err != nil {
- log.Warn("comanage: failed to mark cluster user provisioned", "err", err)
+ log.WarnContext(ctx, "comanage: failed to mark cluster user provisioned", "err", err)
}
return nil
}
@@ -297,18 +297,18 @@ func (o *Orchestrator) findOrCreateUnixClusterGroup(ctx context.Context, cu *mod
span.SetAttributes(attribute.Int("comanage.unix_cluster_group_id", existing))
return nil
} else if err != nil {
- log.Debug("comanage: FindUnixClusterGroup failed; attempting POST", "err", err)
+ log.DebugContext(ctx, "comanage: FindUnixClusterGroup failed; attempting POST", "err", err)
}
if _, err := o.c.CreateUnixClusterGroup(coGroupID); err != nil {
var httpErr *client.HTTPError
if errors.As(err, &httpErr) && httpErr.StatusCode >= 400 && httpErr.StatusCode < 500 {
- log.Info("comanage: UnixClusterGroup attach returned 4xx (already attached)", "status", httpErr.StatusCode)
+ log.InfoContext(ctx, "comanage: UnixClusterGroup attach returned 4xx (already attached)", "status", httpErr.StatusCode)
return nil
}
if errors.As(err, &httpErr) && httpErr.StatusCode == http.StatusInternalServerError {
if existing, gerr := o.c.FindUnixClusterGroup(coGroupID); gerr == nil && existing != 0 {
- log.Info("comanage: UnixClusterGroup POST returned 500 but binding exists", "binding_id", existing)
+ log.InfoContext(ctx, "comanage: UnixClusterGroup POST returned 500 but binding exists", "binding_id", existing)
span.SetAttributes(attribute.Int("comanage.unix_cluster_group_id", existing))
return nil
}
diff --git a/connectors/COmanage/Identity-Provisioner/internal/operations/orchestrator_test.go b/connectors/COmanage/Identity-Provisioner/internal/operations/orchestrator_test.go
index 9d3b107ed..22b42c658 100644
--- a/connectors/COmanage/Identity-Provisioner/internal/operations/orchestrator_test.go
+++ b/connectors/COmanage/Identity-Provisioner/internal/operations/orchestrator_test.go
@@ -26,7 +26,6 @@ import (
"sync"
"testing"
- "go.opentelemetry.io/otel"
sdktrace "go.opentelemetry.io/otel/sdk/trace"
"github.com/apache/airavata-custos/internal/tracing"
@@ -62,11 +61,10 @@ func installRecorder(t *testing.T) *recordingProcessor {
t.Helper()
rec := &recordingProcessor{}
tp := sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(rec))
- prev := otel.GetTracerProvider()
- otel.SetTracerProvider(tp)
+ prev := tracing.SetProvider(tp)
t.Cleanup(func() {
_ = tp.ForceFlush(context.Background())
- otel.SetTracerProvider(prev)
+ tracing.SetProvider(prev)
})
return rec
}
diff --git a/connectors/COmanage/Identity-Provisioner/internal/subscribers/allocation_group.go b/connectors/COmanage/Identity-Provisioner/internal/subscribers/allocation_group.go
index 9f50e07dc..3318d1123 100644
--- a/connectors/COmanage/Identity-Provisioner/internal/subscribers/allocation_group.go
+++ b/connectors/COmanage/Identity-Provisioner/internal/subscribers/allocation_group.go
@@ -21,19 +21,30 @@ import (
"context"
"errors"
+ "github.com/apache/airavata-custos/internal/audit"
+ "github.com/apache/airavata-custos/internal/tracing"
"github.com/apache/airavata-custos/pkg/models"
"github.com/apache/airavata-custos/pkg/service"
)
func (s *ClusterUserSubscriber) handleComputeAllocationChanged(ctx context.Context, a models.ComputeAllocation) error {
+ ctx = audit.WithSource(ctx, "comanage")
+ ctx, span := tracing.Start(ctx, "comanage.compute_allocation_changed")
+ defer span.End()
return s.syncAllocation(ctx, a.ID)
}
func (s *ClusterUserSubscriber) handleMembershipChanged(ctx context.Context, m models.ComputeAllocationMembership) error {
+ ctx = audit.WithSource(ctx, "comanage")
+ ctx, span := tracing.Start(ctx, "comanage.membership_changed")
+ defer span.End()
return s.syncAllocation(ctx, m.ComputeAllocationID)
}
func (s *ClusterUserSubscriber) handleComputeAllocationDeleted(ctx context.Context, a models.ComputeAllocation) error {
+ ctx = audit.WithSource(ctx, "comanage")
+ ctx, span := tracing.Start(ctx, "comanage.compute_allocation_deleted")
+ defer span.End()
if a.ComputeClusterID != s.custosClusterID {
return nil
}
diff --git a/connectors/COmanage/Identity-Provisioner/internal/subscribers/cluster_user.go b/connectors/COmanage/Identity-Provisioner/internal/subscribers/cluster_user.go
index f310396fc..48b936c4a 100644
--- a/connectors/COmanage/Identity-Provisioner/internal/subscribers/cluster_user.go
+++ b/connectors/COmanage/Identity-Provisioner/internal/subscribers/cluster_user.go
@@ -76,8 +76,10 @@ func (s *ClusterUserSubscriber) handleClusterUserDeleted(ctx context.Context, cu
return nil
}
ctx = audit.WithSource(ctx, "comanage")
+ ctx, span := tracing.Start(ctx, "comanage.cluster_user_deleted")
+ defer span.End()
if err := s.ops.RemovePOSIXAccount(ctx, &cu); err != nil {
- slog.Error("comanage subscriber: RemovePOSIXAccount failed", "compute_cluster_user_id", cu.ID, "user_id", cu.UserID, "err", err)
+ slog.ErrorContext(ctx, "comanage subscriber: RemovePOSIXAccount failed", "compute_cluster_user_id", cu.ID, "user_id", cu.UserID, "err", err)
return err
}
return s.syncUserAllocations(ctx, cu.UserID)
@@ -102,17 +104,17 @@ func (s *ClusterUserSubscriber) handleUserIdentityCreate(ctx context.Context, id
cu, err := s.core.GetComputeClusterUserByPair(ctx, s.custosClusterID, identity.UserID)
if errors.Is(err, service.ErrNotFound) {
// Expected for portal admins, system users, and temp users, which have no compute cluster account.
- slog.Debug("comanage subscriber: no cluster account to link the sub to", "user_id", identity.UserID, "cluster_id", s.custosClusterID)
+ slog.DebugContext(ctx, "comanage subscriber: no cluster account to link the sub to", "user_id", identity.UserID, "cluster_id", s.custosClusterID)
return nil
}
if err != nil {
- slog.Error("comanage subscriber: cluster user lookup failed", "user_id", identity.UserID, "err", err)
+ slog.ErrorContext(ctx, "comanage subscriber: cluster user lookup failed", "user_id", identity.UserID, "err", err)
return err
}
if err := s.ops.EnsurePOSIXAccount(ctx, cu); err != nil {
span.RecordError(err)
span.SetStatus(codes.Error, err.Error())
- slog.Error("comanage subscriber: EnsurePOSIXAccount failed after identity link", "compute_cluster_user_id", cu.ID, "user_id", identity.UserID, "err", err)
+ slog.ErrorContext(ctx, "comanage subscriber: EnsurePOSIXAccount failed after identity link", "compute_cluster_user_id", cu.ID, "user_id", identity.UserID, "err", err)
return err
}
return s.syncUserAllocations(ctx, cu.UserID)
@@ -142,7 +144,7 @@ func (s *ClusterUserSubscriber) handleClusterUserApproved(ctx context.Context, c
if err := s.ops.EnsurePOSIXAccount(ctx, &cu); err != nil {
span.RecordError(err)
span.SetStatus(codes.Error, err.Error())
- slog.Error("comanage subscriber: EnsurePOSIXAccount failed", "compute_cluster_user_id", cu.ID, "user_id", cu.UserID, "err", err)
+ slog.ErrorContext(ctx, "comanage subscriber: EnsurePOSIXAccount failed", "compute_cluster_user_id", cu.ID, "user_id", cu.UserID, "err", err)
return err
}
// Even if this fails, the cluster account is usable without sudo, so a
@@ -150,7 +152,7 @@ func (s *ClusterUserSubscriber) handleClusterUserApproved(ctx context.Context, c
if cu.AccessLevel == models.ClusterAccessAdmin {
if err := s.ops.EnsureClusterAdminMembership(ctx, &cu); err != nil {
span.RecordError(err)
- slog.Error("comanage subscriber: EnsureClusterAdminMembership failed", "compute_cluster_user_id", cu.ID, "user_id", cu.UserID, "err", err)
+ slog.ErrorContext(ctx, "comanage subscriber: EnsureClusterAdminMembership failed", "compute_cluster_user_id", cu.ID, "user_id", cu.UserID, "err", err)
return err
}
}
diff --git a/connectors/COmanage/Identity-Provisioner/pkg/comanage/loader.go b/connectors/COmanage/Identity-Provisioner/pkg/comanage/loader.go
index f3746a706..8c290bbd6 100644
--- a/connectors/COmanage/Identity-Provisioner/pkg/comanage/loader.go
+++ b/connectors/COmanage/Identity-Provisioner/pkg/comanage/loader.go
@@ -40,7 +40,7 @@ import (
)
func init() {
- tracing.RegisterTerminalMarkers("comanage", "ComanageClusterAccountAttached")
+ tracing.RegisterMarkers("comanage", "ComanageProvisioningStarted", "ComanageClusterAccountAttached")
}
// Type is the connector type in the config file and the subscriber name on the bus.
diff --git a/connectors/Notification/Email/internal/notify/notifier_integration_test.go b/connectors/Notification/Email/internal/notify/notifier_integration_test.go
index 29c0880f0..947cbfe86 100644
--- a/connectors/Notification/Email/internal/notify/notifier_integration_test.go
+++ b/connectors/Notification/Email/internal/notify/notifier_integration_test.go
@@ -30,6 +30,7 @@ import (
"github.com/google/uuid"
"github.com/apache/airavata-custos/internal/db"
+ "github.com/apache/airavata-custos/internal/tracing/tracingtest"
"github.com/apache/airavata-custos/pkg/events"
"github.com/apache/airavata-custos/pkg/models"
"github.com/apache/airavata-custos/pkg/service"
@@ -70,7 +71,7 @@ func TestNotifier_SendsOneEmailPerAccountEvent(t *testing.T) {
if _, err := database.Exec("TRUNCATE TABLE audit_events, event_deliveries, events, event_subscriptions, compute_cluster_users, compute_clusters, users, organizations CASCADE"); err != nil {
t.Fatal(err)
}
- ctx := context.Background()
+ ctx := tracingtest.Context()
bus, err := events.New(ctx, database)
if err != nil {
t.Fatal(err)
diff --git a/connectors/SLURM/Association-Mapper/internal/subscribers/account.go b/connectors/SLURM/Association-Mapper/internal/subscribers/account.go
index cbb7c3b07..4f4179bc7 100644
--- a/connectors/SLURM/Association-Mapper/internal/subscribers/account.go
+++ b/connectors/SLURM/Association-Mapper/internal/subscribers/account.go
@@ -39,14 +39,14 @@ func (a *AssociationSubscriber) SubscribeToComputeAllocationCreation(ctx context
attribute.String("slurm.cluster_id", computeAllocation.ComputeClusterID),
)
- slog.Info("Received compute allocation creation event", "account", computeAllocation)
+ slog.InfoContext(ctx, "Received compute allocation creation event", "account", computeAllocation)
ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
defer cancel()
cluster, err := a.coreService.GetComputeCluster(ctx, computeAllocation.ComputeClusterID)
if err != nil {
- slog.Error("Failed to get compute cluster for allocation creation", "error", err)
+ slog.ErrorContext(ctx, "Failed to get compute cluster for allocation creation", "error", err)
span.RecordError(err)
span.SetStatus(codes.Error, err.Error())
a.recordAuditEvent(ctx, "ComputeAllocationCreationFailed", "compute_allocation", computeAllocation.ID, "Failed to get compute cluster. Error: "+err.Error())
@@ -55,7 +55,7 @@ func (a *AssociationSubscriber) SubscribeToComputeAllocationCreation(ctx context
project, err := a.coreService.GetProject(ctx, computeAllocation.ProjectID)
if err != nil {
- slog.Error("Failed to get project for allocation creation", "error", err)
+ slog.ErrorContext(ctx, "Failed to get project for allocation creation", "error", err)
span.RecordError(err)
span.SetStatus(codes.Error, err.Error())
a.recordAuditEvent(ctx, "ComputeAllocationCreationFailed", "compute_allocation", computeAllocation.ID, "Failed to get project. Error: "+err.Error())
@@ -64,7 +64,7 @@ func (a *AssociationSubscriber) SubscribeToComputeAllocationCreation(ctx context
pi, err := a.coreService.GetUser(ctx, project.ProjectPIID)
if err != nil {
- slog.Error("Failed to get project PI for allocation creation", "error", err)
+ slog.ErrorContext(ctx, "Failed to get project PI for allocation creation", "error", err)
span.RecordError(err)
span.SetStatus(codes.Error, err.Error())
a.recordAuditEvent(ctx, "ComputeAllocationCreationFailed", "compute_allocation", computeAllocation.ID, "Failed to get project PI. Error: "+err.Error())
@@ -73,7 +73,7 @@ func (a *AssociationSubscriber) SubscribeToComputeAllocationCreation(ctx context
organization, err := a.coreService.GetOrganization(ctx, pi.OrganizationID)
if err != nil {
- slog.Error("Failed to get organization for allocation creation", "error", err)
+ slog.ErrorContext(ctx, "Failed to get organization for allocation creation", "error", err)
span.RecordError(err)
span.SetStatus(codes.Error, err.Error())
a.recordAuditEvent(ctx, "ComputeAllocationCreationFailed", "compute_allocation", computeAllocation.ID, "Failed to get organization. Error: "+err.Error())
@@ -88,7 +88,7 @@ func (a *AssociationSubscriber) SubscribeToComputeAllocationCreation(ctx context
err = a.slurmClient.CreateAccount(slurmAccount, cluster.Name) // TODO: where to get cluster name from?
if err != nil {
- slog.Error("Failed to create SLURM account", "error", err)
+ slog.ErrorContext(ctx, "Failed to create SLURM account", "error", err)
span.RecordError(err)
span.SetStatus(codes.Error, err.Error())
a.recordAuditEvent(ctx, "ComputeAllocationCreationFailed", "compute_allocation", computeAllocation.ID, "Failed to create SLURM account. Error: "+err.Error())
@@ -96,7 +96,7 @@ func (a *AssociationSubscriber) SubscribeToComputeAllocationCreation(ctx context
}
a.recordAuditEvent(ctx, "ComputeAllocationCreationSucceeded", "compute_allocation", computeAllocation.ID, "Successfully created SLURM account for compute allocation.")
- slog.Info("Successfully created SLURM account for compute allocation", "account", slurmAccount)
+ slog.InfoContext(ctx, "Successfully created SLURM account for compute allocation", "account", slurmAccount)
return nil
}
@@ -113,7 +113,7 @@ func (a *AssociationSubscriber) SubscribeToComputeAllocationDeletion(ctx context
defer cancel()
if err := a.removeAssociationsForAllocation(fetchCtx, computeAllocation); err != nil {
- slog.Error("Failed to remove associations for deleted allocation", "error", err)
+ slog.ErrorContext(ctx, "Failed to remove associations for deleted allocation", "error", err)
span.RecordError(err)
span.SetStatus(codes.Error, err.Error())
a.recordAuditEvent(ctx, "ComputeAllocationDeletionFailed", "compute_allocation", computeAllocation.ID, "Failed to remove associations. Error: "+err.Error())
@@ -137,7 +137,7 @@ func (a *AssociationSubscriber) SubscribeToComputeAllocationUpdate(ctx context.C
if activeAllocation(computeAllocation) {
if err := a.restoreAssociationsForAllocation(fetchCtx, computeAllocation); err != nil {
- slog.Error("Failed to restore associations for reactivated allocation", "error", err)
+ slog.ErrorContext(ctx, "Failed to restore associations for reactivated allocation", "error", err)
span.RecordError(err)
span.SetStatus(codes.Error, err.Error())
a.recordAuditEvent(ctx, "ComputeAllocationUpdateFailed", "compute_allocation", computeAllocation.ID, "Failed to restore associations. Error: "+err.Error())
@@ -148,7 +148,7 @@ func (a *AssociationSubscriber) SubscribeToComputeAllocationUpdate(ctx context.C
}
if err := a.removeAssociationsForAllocation(fetchCtx, computeAllocation); err != nil {
- slog.Error("Failed to remove associations for deactivated allocation", "error", err)
+ slog.ErrorContext(ctx, "Failed to remove associations for deactivated allocation", "error", err)
span.RecordError(err)
span.SetStatus(codes.Error, err.Error())
a.recordAuditEvent(ctx, "ComputeAllocationUpdateFailed", "compute_allocation", computeAllocation.ID, "Failed to remove associations. Error: "+err.Error())
@@ -164,14 +164,14 @@ func (a *AssociationSubscriber) SubscribeToComputeAllocationResourceMappingCreat
defer span.End()
span.SetAttributes(attribute.String("slurm.allocation_id", mapping.ComputeAllocationID))
- slog.Info("Received compute allocation resource mapping creation event", "mapping", mapping)
+ slog.InfoContext(ctx, "Received compute allocation resource mapping creation event", "mapping", mapping)
ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
defer cancel()
allocation, err := a.coreService.GetComputeAllocation(ctx, mapping.ComputeAllocationID)
if err != nil {
- slog.Error("Failed to get compute allocation for resource mapping creation", "error", err)
+ slog.ErrorContext(ctx, "Failed to get compute allocation for resource mapping creation", "error", err)
span.RecordError(err)
span.SetStatus(codes.Error, err.Error())
a.recordAuditEvent(ctx, "ComputeAllocationResourceMappingCreationFailed", "compute_allocation_resource_mapping", mapping.ID, "Failed to get compute allocation. Error: "+err.Error())
@@ -180,7 +180,7 @@ func (a *AssociationSubscriber) SubscribeToComputeAllocationResourceMappingCreat
cluster, err := a.coreService.GetComputeCluster(ctx, allocation.ComputeClusterID)
if err != nil {
- slog.Error("Failed to get compute cluster for resource mapping creation", "error", err)
+ slog.ErrorContext(ctx, "Failed to get compute cluster for resource mapping creation", "error", err)
span.RecordError(err)
span.SetStatus(codes.Error, err.Error())
a.recordAuditEvent(ctx, "ComputeAllocationResourceMappingCreationFailed", "compute_allocation_resource_mapping", mapping.ID, "Failed to get compute cluster. Error: "+err.Error())
@@ -190,7 +190,7 @@ func (a *AssociationSubscriber) SubscribeToComputeAllocationResourceMappingCreat
resource, err := a.coreService.GetComputeAllocationResource(ctx, mapping.ComputeAllocationResourceID)
if err != nil {
- slog.Error("Failed to get compute allocation resource for resource mapping creation", "error", err)
+ slog.ErrorContext(ctx, "Failed to get compute allocation resource for resource mapping creation", "error", err)
span.RecordError(err)
span.SetStatus(codes.Error, err.Error())
a.recordAuditEvent(ctx, "ComputeAllocationResourceMappingCreationFailed", "compute_allocation_resource_mapping", mapping.ID, "Failed to get compute allocation resource. Error: "+err.Error())
@@ -202,7 +202,7 @@ func (a *AssociationSubscriber) SubscribeToComputeAllocationResourceMappingCreat
if mapping.ResourceAmount > 0 {
if resource.ResourceType == "" {
- slog.Error("Resource type is empty for resource mapping creation", "mapping", mapping)
+ slog.ErrorContext(ctx, "Resource type is empty for resource mapping creation", "mapping", mapping)
span.SetStatus(codes.Error, "resource type is empty")
a.recordAuditEvent(ctx, "ComputeAllocationResourceMappingCreationFailed", "compute_allocation_resource_mapping", mapping.ID, "Resource type is empty for resource mapping creation")
return errEmptyResourceType
@@ -215,7 +215,7 @@ func (a *AssociationSubscriber) SubscribeToComputeAllocationResourceMappingCreat
if mapping.ResourceTime > 0 {
if resource.ResourceType == "" {
- slog.Error("Resource type is empty for resource mapping creation", "mapping", mapping)
+ slog.ErrorContext(ctx, "Resource type is empty for resource mapping creation", "mapping", mapping)
span.SetStatus(codes.Error, "resource type is empty")
a.recordAuditEvent(ctx, "ComputeAllocationResourceMappingCreationFailed", "compute_allocation_resource_mapping", mapping.ID, "Resource type is empty for resource mapping creation")
return errEmptyResourceType
@@ -237,13 +237,13 @@ func (a *AssociationSubscriber) SubscribeToComputeAllocationResourceMappingCreat
err = a.slurmClient.UpsertAssociation(association)
if err != nil {
- slog.Error("Failed to upsert association for membership resource override creation", "error", err)
+ slog.ErrorContext(ctx, "Failed to upsert association for membership resource override creation", "error", err)
span.RecordError(err)
span.SetStatus(codes.Error, err.Error())
a.recordAuditEvent(ctx, "ComputeAllocationResourceMappingCreationFailed", "compute_allocation_resource_mapping", mapping.ID, "Failed to upsert association. Error: "+err.Error())
return err
} else {
- slog.Info("Successfully upserted association for membership resource override creation", "association", association)
+ slog.InfoContext(ctx, "Successfully upserted association for membership resource override creation", "association", association)
a.recordAuditEvent(ctx, "ComputeAllocationResourceMappingCreationSucceeded", "compute_allocation_resource_mapping", mapping.ID, "Successfully upserted association for compute allocation resource mapping creation")
}
return nil
diff --git a/connectors/SLURM/Association-Mapper/internal/subscribers/association_writer.go b/connectors/SLURM/Association-Mapper/internal/subscribers/association_writer.go
index 400afc8aa..436e5929d 100644
--- a/connectors/SLURM/Association-Mapper/internal/subscribers/association_writer.go
+++ b/connectors/SLURM/Association-Mapper/internal/subscribers/association_writer.go
@@ -44,6 +44,9 @@ func tresFor(resourceType string, count int64) client.TRES {
// and reject the user's jobs. The reconciler retries these.
var errNotProvisioned = errors.New("cluster account not provisioned yet")
+// errNoResources means the allocation has no resource to map onto a partition.
+var errNoResources = errors.New("allocation has no resources")
+
// errEmptyResourceType means the resource row has an empty resource_type. A
// retry cannot fix the row, so the delivery is not retried.
var errEmptyResourceType = fmt.Errorf("%w: resource type is empty", events.ErrPermanent)
@@ -101,6 +104,12 @@ func (a *AssociationSubscriber) upsertAssociationsForMembership(ctx context.Cont
// producing the same record for the same key.
func (a *AssociationSubscriber) syncAssociationsForMembership(ctx context.Context, membership models.ComputeAllocationMembership, existing map[assocKey]client.Association) error {
desired, err := a.desiredAssociationsForMembership(ctx, membership)
+ if errors.Is(err, errNoResources) {
+ // Skipping beats guessing a partition name the cluster may not have.
+ slog.WarnContext(ctx, "Allocation has no resources, no association written",
+ "allocation_id", membership.ComputeAllocationID, "user_id", membership.UserID)
+ return nil
+ }
if err != nil {
return err
}
@@ -111,7 +120,7 @@ func (a *AssociationSubscriber) syncAssociationsForMembership(ctx context.Contex
if err := a.slurmClient.UpsertAssociation(association); err != nil {
return fmt.Errorf("upsert association for partition %s: %w", association.Partition, err)
}
- slog.Info("Upserted association", "association", association)
+ slog.InfoContext(ctx, "Upserted association", "association", association)
}
return nil
}
@@ -144,11 +153,7 @@ func (a *AssociationSubscriber) desiredAssociationsForMembership(ctx context.Con
return nil, fmt.Errorf("list resources for allocation: %w", err)
}
if len(resources) == 0 {
- // Nothing to map onto a partition. Skipping beats guessing a
- // partition name the cluster may not have.
- slog.Warn("Allocation has no resources, no association written",
- "allocation_id", allocation.ID, "user_id", membership.UserID)
- return nil, nil
+ return nil, errNoResources
}
overrides, err := a.coreService.ListOverridesForMembership(ctx, membership.ID)
diff --git a/connectors/SLURM/Association-Mapper/internal/subscribers/deprovision.go b/connectors/SLURM/Association-Mapper/internal/subscribers/deprovision.go
index 98fedb9b6..4cf5a8791 100644
--- a/connectors/SLURM/Association-Mapper/internal/subscribers/deprovision.go
+++ b/connectors/SLURM/Association-Mapper/internal/subscribers/deprovision.go
@@ -64,7 +64,7 @@ func (a *AssociationSubscriber) removeAssociationsForMembership(ctx context.Cont
if err := a.slurmClient.DeleteAssociation(filter); err != nil {
return fmt.Errorf("delete associations for %s on %s: %w", csu.LocalUsername, allocation.Name, err)
}
- slog.Info("Removed associations for membership",
+ slog.InfoContext(ctx, "Removed associations for membership",
"user", csu.LocalUsername, "account", allocation.Name, "cluster", cluster.Name)
return nil
}
@@ -92,7 +92,7 @@ func (a *AssociationSubscriber) removeAssociationsForAllocation(ctx context.Cont
if err := a.slurmClient.DeleteAssociation(filter); err != nil {
return fmt.Errorf("delete associations for allocation %s: %w", allocation.Name, err)
}
- slog.Info("Removed associations for allocation",
+ slog.InfoContext(ctx, "Removed associations for allocation",
"account", allocation.Name, "cluster", cluster.Name)
return nil
}
@@ -121,13 +121,13 @@ func (a *AssociationSubscriber) restoreAssociationsForAllocation(ctx context.Con
continue
}
failed++
- slog.Error("Failed to restore association for member",
+ slog.ErrorContext(ctx, "Failed to restore association for member",
"membership_id", member.ID, "allocation_id", allocation.ID, "error", err)
}
}
if failed > 0 {
return fmt.Errorf("%d of %d members could not be restored", failed, len(members))
}
- slog.Info("Restored associations for allocation", "account", allocation.Name, "members", len(members))
+ slog.InfoContext(ctx, "Restored associations for allocation", "account", allocation.Name, "members", len(members))
return nil
}
diff --git a/connectors/SLURM/Association-Mapper/internal/subscribers/members.go b/connectors/SLURM/Association-Mapper/internal/subscribers/members.go
index 1c9c68b4f..41f15a48d 100644
--- a/connectors/SLURM/Association-Mapper/internal/subscribers/members.go
+++ b/connectors/SLURM/Association-Mapper/internal/subscribers/members.go
@@ -46,7 +46,7 @@ func (a *AssociationSubscriber) SubscribeToComputeAllocationMembershipCreation(c
attribute.String("slurm.user_id", membership.UserID),
)
- slog.Info("Received compute allocation membership creation event", "membership", membership)
+ slog.InfoContext(ctx, "Received compute allocation membership creation event", "membership", membership)
// The short timeout bounds the core fetches only; the provisioning
// wait below runs on its own deadline.
@@ -55,7 +55,7 @@ func (a *AssociationSubscriber) SubscribeToComputeAllocationMembershipCreation(c
allocation, err := a.coreService.GetComputeAllocation(fetchCtx, membership.ComputeAllocationID)
if err != nil {
- slog.Error("Failed to get compute allocation", "error", err)
+ slog.ErrorContext(ctx, "Failed to get compute allocation", "error", err)
span.RecordError(err)
span.SetStatus(codes.Error, err.Error())
a.recordAuditEvent(ctx, "ComputeAllocationMembershipCreationFailed", "compute_allocation_membership", membership.ID, "Failed to get compute allocation. Error: "+err.Error())
@@ -64,7 +64,7 @@ func (a *AssociationSubscriber) SubscribeToComputeAllocationMembershipCreation(c
cluster, err := a.coreService.GetComputeCluster(fetchCtx, allocation.ComputeClusterID)
if err != nil {
- slog.Error("Failed to get compute cluster", "error", err)
+ slog.ErrorContext(ctx, "Failed to get compute cluster", "error", err)
span.RecordError(err)
span.SetStatus(codes.Error, err.Error())
a.recordAuditEvent(ctx, "ComputeAllocationMembershipCreationFailed", "compute_allocation_membership", membership.ID, "Failed to get compute cluster. Error: "+err.Error())
@@ -78,11 +78,11 @@ func (a *AssociationSubscriber) SubscribeToComputeAllocationMembershipCreation(c
if errors.Is(err, errNotProvisioned) {
// Expected on a first-time account: provisioning is still in
// flight. The reconciler writes it once the account lands.
- slog.Info("Cluster account not provisioned yet, leaving the association to the reconciler",
+ slog.InfoContext(ctx, "Cluster account not provisioned yet, leaving the association to the reconciler",
"membership_id", membership.ID, "user_id", membership.UserID)
return nil
}
- slog.Error("Failed to upsert association for membership", "error", err)
+ slog.ErrorContext(ctx, "Failed to upsert association for membership", "error", err)
span.RecordError(err)
span.SetStatus(codes.Error, err.Error())
a.recordAuditEvent(ctx, "ComputeAllocationMembershipCreationFailed", "compute_allocation_membership", membership.ID, "Failed to upsert association. Error: "+err.Error())
@@ -97,14 +97,14 @@ func (a *AssociationSubscriber) SubscribeToComputeAllocationMembershipResourceOv
ctx, span := tracing.Start(ctx, "slurm.compute_allocation_membership_resource_override_create")
defer span.End()
- slog.Info("Received compute allocation membership resource override creation event", "override", override)
+ slog.InfoContext(ctx, "Received compute allocation membership resource override creation event", "override", override)
ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
defer cancel()
// TODO: read per-resource override via
membership, err := a.coreService.GetComputeAllocationMembership(ctx, override.ComputeAllocationMembershipID)
if err != nil {
- slog.Error("Failed to get compute allocation membership for resource override creation", "error", err)
+ slog.ErrorContext(ctx, "Failed to get compute allocation membership for resource override creation", "error", err)
span.RecordError(err)
span.SetStatus(codes.Error, err.Error())
a.recordAuditEvent(ctx, "ComputeAllocationMembershipResourceOverrideCreationFailed", "compute_allocation_membership_resource_override", override.ID, "Failed to get compute allocation membership. Error: "+err.Error())
@@ -120,11 +120,11 @@ func (a *AssociationSubscriber) SubscribeToComputeAllocationMembershipResourceOv
// than only this override's slice.
if err := a.upsertAssociationsForMembership(ctx, *membership); err != nil {
if errors.Is(err, errNotProvisioned) {
- slog.Info("Cluster account not provisioned yet, leaving the association to the reconciler",
+ slog.InfoContext(ctx, "Cluster account not provisioned yet, leaving the association to the reconciler",
"membership_id", membership.ID, "override_id", override.ID)
return nil
}
- slog.Error("Failed to upsert association for membership resource override", "error", err)
+ slog.ErrorContext(ctx, "Failed to upsert association for membership resource override", "error", err)
span.RecordError(err)
span.SetStatus(codes.Error, err.Error())
a.recordAuditEvent(ctx, "ComputeAllocationMembershipResourceOverrideCreationFailed", "compute_allocation_membership_resource_override", override.ID, "Failed to upsert association. Error: "+err.Error())
@@ -152,11 +152,11 @@ func (a *AssociationSubscriber) SubscribeToComputeAllocationMembershipUpdate(ctx
if membership.MembershipStatus == models.ACTIVE {
if err := a.upsertAssociationsForMembership(fetchCtx, membership); err != nil {
if errors.Is(err, errNotProvisioned) {
- slog.Info("Cluster account not provisioned yet, leaving the association to the reconciler",
+ slog.InfoContext(ctx, "Cluster account not provisioned yet, leaving the association to the reconciler",
"membership_id", membership.ID)
return nil
}
- slog.Error("Failed to upsert association for reactivated membership", "error", err)
+ slog.ErrorContext(ctx, "Failed to upsert association for reactivated membership", "error", err)
span.RecordError(err)
span.SetStatus(codes.Error, err.Error())
a.recordAuditEvent(ctx, "ComputeAllocationMembershipUpdateFailed", "compute_allocation_membership", membership.ID, "Failed to upsert association. Error: "+err.Error())
@@ -166,7 +166,7 @@ func (a *AssociationSubscriber) SubscribeToComputeAllocationMembershipUpdate(ctx
}
if err := a.removeAssociationsForMembership(fetchCtx, membership); err != nil {
- slog.Error("Failed to remove associations for deactivated membership", "error", err)
+ slog.ErrorContext(ctx, "Failed to remove associations for deactivated membership", "error", err)
span.RecordError(err)
span.SetStatus(codes.Error, err.Error())
a.recordAuditEvent(ctx, "ComputeAllocationMembershipUpdateFailed", "compute_allocation_membership", membership.ID, "Failed to remove associations. Error: "+err.Error())
@@ -191,7 +191,7 @@ func (a *AssociationSubscriber) SubscribeToComputeAllocationMembershipDeletion(c
defer cancel()
if err := a.removeAssociationsForMembership(fetchCtx, membership); err != nil {
- slog.Error("Failed to remove associations for deleted membership", "error", err)
+ slog.ErrorContext(ctx, "Failed to remove associations for deleted membership", "error", err)
span.RecordError(err)
span.SetStatus(codes.Error, err.Error())
a.recordAuditEvent(ctx, "ComputeAllocationMembershipDeletionFailed", "compute_allocation_membership", membership.ID, "Failed to remove associations. Error: "+err.Error())
diff --git a/connectors/SLURM/Association-Mapper/internal/subscribers/reconciler.go b/connectors/SLURM/Association-Mapper/internal/subscribers/reconciler.go
index c761a4bac..a5098492f 100644
--- a/connectors/SLURM/Association-Mapper/internal/subscribers/reconciler.go
+++ b/connectors/SLURM/Association-Mapper/internal/subscribers/reconciler.go
@@ -19,10 +19,12 @@ package subscribers
import (
"context"
+ "errors"
"log/slog"
"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,12 +77,14 @@ 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()
clusters, err := a.coreService.ListComputeClusters(sweepCtx)
if err != nil {
- slog.Error("Association reconciler: failed to list compute clusters", "error", err)
+ slog.ErrorContext(ctx, "Association reconciler: failed to list compute clusters", "error", err)
return
}
@@ -92,7 +96,7 @@ func (a *AssociationSubscriber) reconcile(ctx context.Context) {
func (a *AssociationSubscriber) reconcileCluster(ctx context.Context, cluster models.ComputeCluster) {
clusterUsers, err := a.coreService.ListComputeClusterUsersByCluster(ctx, cluster.ID)
if err != nil {
- slog.Error("Association reconciler: failed to list cluster users", "cluster_id", cluster.ID, "error", err)
+ slog.ErrorContext(ctx, "Association reconciler: failed to list cluster users", "cluster_id", cluster.ID, "error", err)
return
}
@@ -100,7 +104,7 @@ func (a *AssociationSubscriber) reconcileCluster(ctx context.Context, cluster mo
// association, which costs extra calls but still converges.
existing, err := a.existingAssociations(ctx, cluster.Name)
if err != nil {
- slog.Warn("Association reconciler: could not read current associations, writing unconditionally",
+ slog.WarnContext(ctx, "Association reconciler: could not read current associations, writing unconditionally",
"cluster", cluster.Name, "error", err)
existing = nil
}
@@ -113,7 +117,7 @@ func (a *AssociationSubscriber) reconcileCluster(ctx context.Context, cluster mo
}
memberships, err := a.coreService.ListAllocationsForUser(ctx, csu.UserID)
if err != nil {
- slog.Error("Association reconciler: failed to list memberships", "user_id", csu.UserID, "error", err)
+ slog.ErrorContext(ctx, "Association reconciler: failed to list memberships", "user_id", csu.UserID, "error", err)
// Without this user's memberships the desired set is incomplete,
// so pruning could revoke access that is actually valid.
return
@@ -123,8 +127,12 @@ func (a *AssociationSubscriber) reconcileCluster(ctx context.Context, cluster mo
continue
}
records, err := a.desiredAssociationsForMembership(ctx, membership)
+ // Already warned when the membership was written; nothing to sweep.
+ if errors.Is(err, errNoResources) {
+ continue
+ }
if err != nil {
- slog.Error("Association reconciler: failed to resolve desired associations",
+ slog.ErrorContext(ctx, "Association reconciler: failed to resolve desired associations",
"membership_id", membership.ID, "error", err)
return
}
@@ -141,7 +149,7 @@ func (a *AssociationSubscriber) reconcileCluster(ctx context.Context, cluster mo
continue
}
if err := a.slurmClient.UpsertAssociation(record); err != nil {
- slog.Error("Association reconciler: failed to upsert association",
+ slog.ErrorContext(ctx, "Association reconciler: failed to upsert association",
"membership_id", membership.ID, "error", err)
continue
}
@@ -151,7 +159,7 @@ func (a *AssociationSubscriber) reconcileCluster(ctx context.Context, cluster mo
}
removed := a.pruneStaleAssociations(ctx, cluster, existing, desired)
- slog.Debug("Association reconciler pass complete",
+ slog.DebugContext(ctx, "Association reconciler pass complete",
"cluster", cluster.Name, "written", written, "removed", removed)
}
@@ -167,14 +175,14 @@ func (a *AssociationSubscriber) pruneStaleAssociations(ctx context.Context, clus
// nobody being entitled to anything, and acting on it would revoke the
// whole cluster. Refuse.
if len(desired) == 0 {
- slog.Warn("Association reconciler: nothing is desired on this cluster, skipping removal",
+ slog.WarnContext(ctx, "Association reconciler: nothing is desired on this cluster, skipping removal",
"cluster", cluster.Name, "existing", len(existing))
return 0
}
managed, err := a.managedAccounts(ctx, cluster.ID)
if err != nil {
- slog.Error("Association reconciler: could not list allocations, skipping removal",
+ slog.ErrorContext(ctx, "Association reconciler: could not list allocations, skipping removal",
"cluster", cluster.Name, "error", err)
return 0
}
@@ -197,11 +205,11 @@ func (a *AssociationSubscriber) pruneStaleAssociations(ctx context.Context, clus
Partition: assoc.Partition,
}
if err := a.slurmClient.DeleteAssociation(filter); err != nil {
- slog.Error("Association reconciler: failed to remove stale association",
+ slog.ErrorContext(ctx, "Association reconciler: failed to remove stale association",
"account", assoc.Account, "user", assoc.User, "partition", assoc.Partition, "error", err)
continue
}
- slog.Info("Removed stale association",
+ slog.InfoContext(ctx, "Removed stale association",
"account", assoc.Account, "user", assoc.User, "partition", assoc.Partition)
removed++
}
diff --git a/connectors/SLURM/Association-Mapper/internal/subscribers/subscriber_trace_test.go b/connectors/SLURM/Association-Mapper/internal/subscribers/subscriber_trace_test.go
index a680315e1..ce083aa19 100644
--- a/connectors/SLURM/Association-Mapper/internal/subscribers/subscriber_trace_test.go
+++ b/connectors/SLURM/Association-Mapper/internal/subscribers/subscriber_trace_test.go
@@ -25,18 +25,16 @@ import (
"github.com/apache/airavata-custos/internal/tracing"
"github.com/apache/airavata-custos/pkg/models"
"github.com/apache/airavata-custos/pkg/service"
- "go.opentelemetry.io/otel"
sdktrace "go.opentelemetry.io/otel/sdk/trace"
"go.opentelemetry.io/otel/sdk/trace/tracetest"
)
func setupRecorder(t *testing.T) *tracetest.SpanRecorder {
t.Helper()
- prev := otel.GetTracerProvider()
sr := tracetest.NewSpanRecorder()
tp := sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(sr))
- otel.SetTracerProvider(tp)
- t.Cleanup(func() { otel.SetTracerProvider(prev) })
+ prev := tracing.SetProvider(tp)
+ t.Cleanup(func() { tracing.SetProvider(prev) })
return sr
}
diff --git a/connectors/SLURM/Usage-Monitor/internal/smonitor/smonitor.go b/connectors/SLURM/Usage-Monitor/internal/smonitor/smonitor.go
index 270dd24fb..69116ed3a 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,17 +106,19 @@ func (m *SlurmMonitor) StartMonitor(ctx context.Context) {
func (m *SlurmMonitor) poll(ctx context.Context) {
slog.Debug("polling SLURM usage")
+ ctx, span := tracing.Start(ctx, "slurm.usage_poll")
+ defer span.End()
ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
defer cancel()
cluster, err := m.coreService.GetComputeCluster(ctx, m.clusterId)
if err != nil {
- slog.Error("failed to get compute cluster", "error", err)
+ slog.ErrorContext(ctx, "failed to get compute cluster", "error", err)
return
}
allocations, err := m.coreService.ListComputeAllocationsByCluster(ctx, cluster.ID)
if err != nil {
- slog.Error("failed to list compute allocations", "error", err)
+ slog.ErrorContext(ctx, "failed to list compute allocations", "error", err)
return
}
@@ -126,7 +129,7 @@ func (m *SlurmMonitor) poll(ctx context.Context) {
jobs, err := m.slurmClient.ListJobs(jobFilter)
if err != nil {
- slog.Error("failed to list SLURM jobs", "error", err)
+ slog.ErrorContext(ctx, "failed to list SLURM jobs", "error", err)
return
}
m.lastMonitorTime = jobFilter.EndTime
@@ -135,27 +138,27 @@ func (m *SlurmMonitor) poll(ctx context.Context) {
m.recordJob(ctx, job, cluster, allocations)
}
- slog.Info("successfully polled SLURM usage", "num_allocations", len(allocations), "num_jobs", len(jobs))
+ slog.InfoContext(ctx, "successfully polled SLURM usage", "num_allocations", len(allocations), "num_jobs", len(jobs))
}
// recordJob writes one usage row for a matched job; returning skips only this
// job so one bad lookup cannot abort the rest of the poll cycle.
func (m *SlurmMonitor) recordJob(ctx context.Context, job client.JobInfo, cluster *models.ComputeCluster, allocations []models.ComputeAllocation) {
- slog.Info("Job object", "job", job)
+ slog.InfoContext(ctx, "Job object", "job", job)
targetAccount := job.Account
for _, alloc := range allocations {
if alloc.Name != targetAccount {
continue
}
- slog.Info("found matching compute allocation for SLURM job", "job_id", job.JobID, "allocation_id", alloc.ID)
+ slog.InfoContext(ctx, "found matching compute allocation for SLURM job", "job_id", job.JobID, "allocation_id", alloc.ID)
user, err := m.coreService.GetComputeClusterUserByClusterAndLocalUsername(ctx, cluster.ID, job.User)
if err != nil {
if err == service.ErrNotFound {
- slog.Warn("compute cluster user not found for SLURM job, skipping usage recording", "local_username", job.User, "cluster_id", cluster.ID)
+ slog.WarnContext(ctx, "compute cluster user not found for SLURM job, skipping usage recording", "local_username", job.User, "cluster_id", cluster.ID)
} else {
- slog.Error("failed to get compute cluster user", "error", err)
+ slog.ErrorContext(ctx, "failed to get compute cluster user", "error", err)
}
return
}
@@ -163,9 +166,9 @@ func (m *SlurmMonitor) recordJob(ctx context.Context, job client.JobInfo, cluste
resource, err := m.coreService.GetComputeAllocationResourceByNameAndCluster(ctx, job.Partition, cluster.ID)
if err != nil {
if err == service.ErrNotFound {
- slog.Warn("compute allocation resource not found for SLURM job, skipping usage recording", "resource_name", job.Partition, "cluster_id", cluster.ID)
+ slog.WarnContext(ctx, "compute allocation resource not found for SLURM job, skipping usage recording", "resource_name", job.Partition, "cluster_id", cluster.ID)
} else {
- slog.Error("failed to get compute allocation resource", "error", err)
+ slog.ErrorContext(ctx, "failed to get compute allocation resource", "error", err)
}
return
}
@@ -173,13 +176,13 @@ func (m *SlurmMonitor) recordJob(ctx context.Context, job client.JobInfo, cluste
jobId := strconv.FormatInt(job.JobID, 10)
existing, err := m.coreService.GetComputeAllocationUsageByComputeAllocationIDAndJobID(ctx, alloc.ID, jobId)
if err != nil && err != service.ErrNotFound {
- slog.Error("failed to check for existing compute allocation usage", "error", err)
+ slog.ErrorContext(ctx, "failed to check for existing compute allocation usage", "error", err)
return
}
jobDurationSec := job.Time.End - job.Time.Start
if jobDurationSec <= 0 {
- slog.Warn("SLURM job has non-positive duration, skipping usage recording", "job_id", job.JobID, "duration_seconds", jobDurationSec)
+ slog.WarnContext(ctx, "SLURM job has non-positive duration, skipping usage recording", "job_id", job.JobID, "duration_seconds", jobDurationSec)
return
}
@@ -192,7 +195,7 @@ func (m *SlurmMonitor) recordJob(ctx context.Context, job client.JobInfo, cluste
}
}
if billingAmount <= 0 {
- slog.Warn("SLURM job has no billing charge, skipping usage recording", "job_id", job.JobID, "billing", billingAmount)
+ slog.WarnContext(ctx, "SLURM job has no billing charge, skipping usage recording", "job_id", job.JobID, "billing", billingAmount)
return
}
@@ -202,9 +205,9 @@ func (m *SlurmMonitor) recordJob(ctx context.Context, job client.JobInfo, cluste
rate, err := m.coreService.GetEffectiveRateForResource(ctx, resource.ID, time.Unix(job.Time.End, 0))
if err != nil {
if err == service.ErrNotFound {
- slog.Warn("no rate covers the job end time, skipping usage recording", "job_id", job.JobID, "resource_id", resource.ID)
+ slog.WarnContext(ctx, "no rate covers the job end time, skipping usage recording", "job_id", job.JobID, "resource_id", resource.ID)
} else {
- slog.Error("failed to get effective rate for resource", "error", err, "job_id", job.JobID, "resource_id", resource.ID)
+ slog.ErrorContext(ctx, "failed to get effective rate for resource", "error", err, "job_id", job.JobID, "resource_id", resource.ID)
}
return
}
@@ -221,7 +224,7 @@ func (m *SlurmMonitor) recordJob(ctx context.Context, job client.JobInfo, cluste
if existing != nil {
m.coreService.DeleteComputeAllocationUsage(ctx, existing.ID)
- slog.Info("deleted existing compute allocation usage for SLURM job", "job_id", job.JobID, "existing_usage_id", existing.ID)
+ slog.InfoContext(ctx, "deleted existing compute allocation usage for SLURM job", "job_id", job.JobID, "existing_usage_id", existing.ID)
}
m.coreService.CreateComputeAllocationUsage(ctx, usageModel)
return
diff --git a/connectors/Storage/VAST-Provisioner/internal/subscribers/subscriber.go b/connectors/Storage/VAST-Provisioner/internal/subscribers/subscriber.go
index 407a1f8a9..e1913e8d4 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 {
@@ -179,6 +182,6 @@ func (s *StorageSubscriber) reconcile(ctx context.Context, id string, c Cluster)
errs = append(errs, s.ensureAllocation(ctx, a))
}
if err := errors.Join(errs...); err != nil {
- slog.Error("VAST reconciler", "cluster_id", id, "err", err)
+ slog.ErrorContext(ctx, "VAST reconciler", "cluster_id", id, "err", err)
}
}
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..6e0c79931 100644
--- a/internal/db/migrations/000007_event_bus.up.sql
+++ b/internal/db/migrations/000007_event_bus.up.sql
@@ -35,7 +35,8 @@ 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 <> ''),
+ span_id CHAR(16) NOT NULL CHECK (span_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/cluster_account_integration_test.go b/internal/server/cluster_account_integration_test.go
index 6237a4c57..159a6bb00 100644
--- a/internal/server/cluster_account_integration_test.go
+++ b/internal/server/cluster_account_integration_test.go
@@ -29,6 +29,7 @@ import (
"github.com/google/uuid"
"github.com/jmoiron/sqlx"
+ "github.com/apache/airavata-custos/internal/tracing/tracingtest"
"github.com/apache/airavata-custos/pkg/models"
"github.com/apache/airavata-custos/pkg/service"
)
@@ -59,7 +60,7 @@ func TestListClusterAccounts_FiltersByApprovalStatus(t *testing.T) {
admin := seedUser(t, database, fmt.Sprintf("admin-%s@example.edu", uuid.NewString()))
pending := seedClusterUser(t, database, svc)
approved := seedClusterUser(t, database, svc)
- if _, err := svc.ApproveComputeClusterUser(t.Context(), approved.ID, admin); err != nil {
+ if _, err := svc.ApproveComputeClusterUser(tracingtest.Context(), approved.ID, admin); err != nil {
t.Fatalf("approve: %v", err)
}
diff --git a/internal/server/event_delivery_integration_test.go b/internal/server/event_delivery_integration_test.go
index 834dd4666..f35972318 100644
--- a/internal/server/event_delivery_integration_test.go
+++ b/internal/server/event_delivery_integration_test.go
@@ -28,6 +28,7 @@ import (
"testing"
"time"
+ "github.com/apache/airavata-custos/internal/tracing/tracingtest"
"github.com/apache/airavata-custos/pkg/events"
"github.com/apache/airavata-custos/pkg/models"
"github.com/apache/airavata-custos/pkg/service"
@@ -72,7 +73,7 @@ func failDelivery(t *testing.T, svc *service.Service, bus *events.Bus) models.Pe
}
return nil
})
- if _, err := svc.CreateOrganization(context.Background(), &models.Organization{OriginatedID: "org-retry", Name: "Retry Org"}); err != nil {
+ if _, err := svc.CreateOrganization(tracingtest.Context(), &models.Organization{OriginatedID: "org-retry", Name: "Retry Org"}); err != nil {
t.Fatalf("create organization: %v", err)
}
return runWorkerUntil(t, bus, models.EventDeliveryFailed)
@@ -167,7 +168,7 @@ func TestRetryRefusesDeliveryThatHasNotFailed(t *testing.T) {
_, svc, srv := setupTestStack(t)
bus := svc.EventBus()
bus.Subscribe(subscriberTest, events.OrganizationCreateEvent, func(context.Context, events.Event, any) error { return nil })
- if _, err := svc.CreateOrganization(context.Background(), &models.Organization{OriginatedID: "org-ok", Name: "Ok Org"}); err != nil {
+ if _, err := svc.CreateOrganization(tracingtest.Context(), &models.Organization{OriginatedID: "org-ok", Name: "Ok Org"}); err != nil {
t.Fatalf("create organization: %v", err)
}
done := runWorkerUntil(t, bus, models.EventDeliverySucceeded)
diff --git a/internal/server/privilege_integration_test.go b/internal/server/privilege_integration_test.go
index 87f4d5d85..5d1c90a90 100644
--- a/internal/server/privilege_integration_test.go
+++ b/internal/server/privilege_integration_test.go
@@ -27,6 +27,8 @@ import (
"net/http/httptest"
"testing"
+ "github.com/apache/airavata-custos/internal/tracing"
+ "github.com/apache/airavata-custos/internal/tracing/tracingtest"
"github.com/apache/airavata-custos/pkg/models"
)
@@ -214,10 +216,14 @@ func TestGrantPrivilegeEndpoint_HappyPath(t *testing.T) {
req := httptest.NewRequest(http.MethodPost, "/users/"+target+"/privileges", bytes.NewReader(body))
req.Header.Set("Content-Type", "application/json")
req = withTestCaller(req, granter, models.PrivilegesGrant)
- srv.ServeHTTP(rr, req)
+ tracing.Middleware(srv).ServeHTTP(rr, req)
if rr.Code != http.StatusCreated {
t.Fatalf("status: got %d, want 201, body=%s", rr.Code, rr.Body.String())
}
+ var traceID string
+ if err := database.Get(&traceID, "SELECT trace_id FROM audit_events WHERE entity_id = $1", target); err != nil || traceID != rr.Header().Get("X-Trace-Id") {
+ t.Fatalf("audit row trace=%q err=%v, want the request's trace", traceID, err)
+ }
has, err := svc.HasPrivilege(context.Background(), target, models.ClustersRead)
if err != nil || !has {
t.Errorf("HasPrivilege after grant endpoint: has=%v err=%v", has, err)
@@ -244,7 +250,7 @@ func TestRevokePrivilegeEndpoint_HappyPath(t *testing.T) {
granter := seedUser(t, database, "granter@example.edu")
target := seedUser(t, database, "target@example.edu")
seedPrivilegesGrant(t, database, granter)
- if _, err := svc.GrantPrivilege(context.Background(), target, models.ClustersRead, granter, ""); err != nil {
+ if _, err := svc.GrantPrivilege(tracingtest.Context(), target, models.ClustersRead, granter, ""); err != nil {
t.Fatalf("seed grant: %v", err)
}
body, _ := json.Marshal(map[string]any{"reason": "rotated"})
@@ -278,7 +284,7 @@ func TestListUserPrivilegesEndpoint(t *testing.T) {
granter := seedUser(t, database, "granter@example.edu")
target := seedUser(t, database, "target@example.edu")
seedPrivilegesGrant(t, database, granter)
- if _, err := svc.GrantPrivilege(context.Background(), target, models.ClustersRead, granter, ""); err != nil {
+ if _, err := svc.GrantPrivilege(tracingtest.Context(), target, models.ClustersRead, granter, ""); err != nil {
t.Fatalf("grant: %v", err)
}
rr := httptest.NewRecorder()
@@ -302,7 +308,7 @@ func TestListPrivilegeHoldersEndpoint(t *testing.T) {
granter := seedUser(t, database, "granter@example.edu")
target := seedUser(t, database, "target@example.edu")
seedPrivilegesGrant(t, database, granter)
- if _, err := svc.GrantPrivilege(context.Background(), target, models.ClustersRead, granter, ""); err != nil {
+ if _, err := svc.GrantPrivilege(tracingtest.Context(), target, models.ClustersRead, granter, ""); err != nil {
t.Fatalf("grant: %v", err)
}
rr := httptest.NewRecorder()
diff --git a/internal/server/server.go b/internal/server/server.go
index 6ed954037..8513c2cc6 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"
@@ -48,6 +49,7 @@ func New(svc *service.Service, router *identity.Router) *Server {
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
s.router.ServeHTTP(w, r)
+ tracing.SetRoute(r)
}
// requireCaller pulls the verified caller off the request context. Returns
@@ -207,25 +209,21 @@ type statusUpdateRequest struct {
Status string `json:"status"`
}
-// LoggingMiddleware logs every request once it completes. It wraps
-// tracing.Middleware (logging outer, tracing inner) and reads the trace_id
-// 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.
+// LoggingMiddleware logs every request once it completes. It sits outside
+// tracing.Middleware and reads the trace_id from the X-Trace-Id response
+// header, since the inner request ctx is not visible here after ServeHTTP returns.
func LoggingMiddleware(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
start := time.Now()
rw := &httputil.StatusRecorder{ResponseWriter: w, Status: http.StatusOK}
next.ServeHTTP(rw, r)
- attrs := []any{
+ slog.InfoContext(r.Context(), "http request",
"method", r.Method,
"path", r.URL.Path,
"status", rw.Status,
"duration", time.Since(start).String(),
- }
- if tid := rw.Header().Get("X-Trace-Id"); tid != "" {
- attrs = append(attrs, "trace_id", tid)
- }
- slog.InfoContext(r.Context(), "http request", attrs...)
+ "trace_id", rw.Header().Get("X-Trace-Id"),
+ )
})
}
diff --git a/internal/server/server_trace_test.go b/internal/server/server_trace_test.go
index 40f86525c..51e3df092 100644
--- a/internal/server/server_trace_test.go
+++ b/internal/server/server_trace_test.go
@@ -18,30 +18,31 @@
package server
import (
+ "context"
"net/http"
"net/http/httptest"
"testing"
- "go.opentelemetry.io/otel"
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"
+ "github.com/apache/airavata-custos/pkg/models"
)
// TestLoggingWrapsTracingProducesTraceIdHeader verifies the cmd/server stack
-// composition: LoggingMiddleware(tracing.Middleware(server.New(...))) lets the
+// composition: LoggingMiddleware(tracing.Middleware(identity.Middleware(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()
sr := tracetest.NewSpanRecorder()
tp := sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(sr))
- otel.SetTracerProvider(tp)
- t.Cleanup(func() { otel.SetTracerProvider(prev) })
+ prev := tracing.SetProvider(tp)
+ t.Cleanup(func() { tracing.SetProvider(prev) })
router := identity.NewRouter(http.NewServeMux())
- handler := LoggingMiddleware(tracing.Middleware(New(nil, router)))
+ srv := New(nil, router)
+ handler := LoggingMiddleware(tracing.Middleware(identity.Middleware(stubAuth{}, stubAuth{}, router.PublicPaths(), srv)))
req := httptest.NewRequest(http.MethodGet, "/healthz", nil)
rec := httptest.NewRecorder()
@@ -55,11 +56,29 @@ func TestLoggingWrapsTracingProducesTraceIdHeader(t *testing.T) {
t.Fatalf("expected 32-char X-Trace-Id header, got %q", hdr)
}
+ // Authenticated route: the span takes the mux pattern, not the raw path.
+ req = httptest.NewRequest(http.MethodGet, "/users/abc", nil)
+ req.Header.Set("Authorization", "Bearer t")
+ handler.ServeHTTP(httptest.NewRecorder(), req)
+
spans := sr.Ended()
- if len(spans) != 1 {
- t.Fatalf("expected exactly 1 span, got %d", len(spans))
+ if len(spans) != 2 {
+ t.Fatalf("expected exactly 2 spans, got %d", len(spans))
}
if got, want := spans[0].Name(), "http.GET /healthz"; got != want {
t.Fatalf("span name = %q, want %q", got, want)
}
+ if got, want := spans[1].Name(), "http.GET /users/{id}"; got != want {
+ t.Fatalf("span name = %q, want %q", got, want)
+ }
+}
+
+type stubAuth struct{}
+
+func (stubAuth) Verify(context.Context, string) (*identity.Claims, error) {
+ return &identity.Claims{}, nil
+}
+
+func (stubAuth) ResolveCaller(context.Context, *identity.Claims) (*identity.Caller, []models.PrivilegeKey, error) {
+ return &identity.Caller{}, nil, nil
}
diff --git a/internal/server/testidentity_test.go b/internal/server/testidentity_test.go
index 4e98db01e..a840a5f5b 100644
--- a/internal/server/testidentity_test.go
+++ b/internal/server/testidentity_test.go
@@ -20,6 +20,7 @@ package server
import (
"net/http"
+ "github.com/apache/airavata-custos/internal/tracing"
"github.com/apache/airavata-custos/pkg/identity"
"github.com/apache/airavata-custos/pkg/models"
)
@@ -27,7 +28,8 @@ import (
// withTestCaller installs a *identity.Caller plus a privilege set onto the
// request context for handler tests that run without the auth middleware.
func withTestCaller(r *http.Request, userID string, privs ...models.PrivilegeKey) *http.Request {
- ctx := identity.WithCaller(r.Context(), &identity.Caller{UserID: userID})
+ ctx, _ := tracing.Start(r.Context(), "test.request")
+ ctx = identity.WithCaller(ctx, &identity.Caller{UserID: userID})
ctx = identity.WithPrivilegesForTest(ctx, privs)
return r.WithContext(ctx)
}
diff --git a/internal/store/audit_event_store.go b/internal/store/audit_event_store.go
index fb849dd87..9df4d555f 100644
--- a/internal/store/audit_event_store.go
+++ b/internal/store/audit_event_store.go
@@ -18,12 +18,16 @@
package store
import (
+ "cmp"
"context"
"database/sql"
"errors"
+ "fmt"
"github.com/jmoiron/sqlx"
+ "github.com/apache/airavata-custos/internal/audit"
+ "github.com/apache/airavata-custos/internal/tracing"
"github.com/apache/airavata-custos/pkg/models"
)
@@ -90,6 +94,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 err := tracing.PopulateAuditIDs(ctx, &e.TraceID, &e.SpanID, &e.ParentSpanID); err != nil {
+ return fmt.Errorf("audit %s: %w", e.EventType, err)
+ }
+ e.Source = cmp.Or(e.Source, audit.SourceFromContext(ctx), "core")
_, 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..d2cc09f66 100644
--- a/internal/store/audit_trace_store.go
+++ b/internal/store/audit_trace_store.go
@@ -20,6 +20,7 @@ package store
import (
"context"
"fmt"
+ "slices"
"strings"
"time"
@@ -40,6 +41,8 @@ type TraceFilter struct {
Q string
Limit int
Offset int
+
+ traceIDs []string // set by ListTraces once a status filter has picked the traces
}
type AuditTraceStore interface {
@@ -66,7 +69,6 @@ SELECT
details AS description,
event_time AS created_at
FROM audit_events
-WHERE trace_id <> ''
`
type pgAuditTraceStore struct {
@@ -80,7 +82,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"`
@@ -103,9 +105,23 @@ func (r rowEvent) toTraceEvent() models.TraceEvent {
}
}
-// ListTraces aggregates by trace_id then runs one summary probe per row.
-// Acceptable while page sizes stay <= 200.
+// ListTraces pages traces in SQL and summarises the page in one more query.
+// Status is computed in Go, so a status filter first loads every matching
+// trace's rows in one query and narrows the trace set before counting and paging.
func (s *pgAuditTraceStore) ListTraces(ctx context.Context, f TraceFilter) ([]models.TraceSummary, int, error) {
+ if len(f.Statuses) > 0 {
+ whereSQL, args := buildTraceWhere(f)
+ byTrace, err := s.traceRows(ctx, `trace_id IN (SELECT trace_id FROM (`+auditTraceSelect+`) u `+whereSQL+`)`, args...)
+ if err != nil {
+ return nil, 0, err
+ }
+ f.traceIDs = []string{}
+ for id, rows := range byTrace {
+ if _, _, status := summarise(rows); slices.Contains(f.Statuses, status) {
+ f.traceIDs = append(f.traceIDs, id)
+ }
+ }
+ }
whereSQL, args := buildTraceWhere(f)
limit := f.Limit
@@ -144,15 +160,17 @@ func (s *pgAuditTraceStore) ListTraces(ctx context.Context, f TraceFilter) ([]mo
return nil, 0, fmt.Errorf("audit_trace_store: list: %w", err)
}
+ ids := make([]string, len(raw))
+ for i, r := range raw {
+ ids[i] = r.TraceID
+ }
+ byTrace, err := s.traceRows(ctx, `trace_id = ANY(?)`, ids)
+ if err != nil {
+ return nil, 0, err
+ }
out := make([]models.TraceSummary, 0, len(raw))
for _, r := range raw {
- rootOp, src, status, err := s.summariseTrace(ctx, r.TraceID)
- if err != nil {
- return nil, 0, err
- }
- if !matchesStatusFilter(status, f.Statuses) {
- continue
- }
+ rootOp, src, status := summarise(byTrace[r.TraceID])
out = append(out, models.TraceSummary{
TraceID: r.TraceID,
RootOperation: rootOp,
@@ -166,80 +184,70 @@ func (s *pgAuditTraceStore) ListTraces(ctx context.Context, f TraceFilter) ([]mo
return out, total, nil
}
-func (s *pgAuditTraceStore) summariseTrace(ctx context.Context, traceID string) (string, string, string, error) {
+// traceRows loads the rows of every trace matching where, oldest first, keyed by trace.
+func (s *pgAuditTraceStore) traceRows(ctx context.Context, where string, args ...any) (map[string][]rowEvent, error) {
q := `SELECT trace_id, span_id, parent_span_id, source, event_type, entity_type, entity_id, description, created_at
FROM (` + auditTraceSelect + `) u
- WHERE trace_id = ?
+ WHERE ` + where + `
ORDER BY created_at ASC, span_id ASC`
var rows []rowEvent
- if err := s.db.SelectContext(ctx, &rows, s.db.Rebind(q), traceID); err != nil {
- return "", "", "", fmt.Errorf("audit_trace_store: summarise: %w", err)
- }
- if len(rows) == 0 {
- return "", "", tracing.StatusInProgress, nil
- }
-
- var rootOp string
- for _, r := range rows {
- if r.ParentSpanID == "" {
- rootOp = r.EventType
- break
- }
- }
- if rootOp == "" {
- rootOp = rows[0].EventType
+ if err := s.db.SelectContext(ctx, &rows, s.db.Rebind(q), args...); err != nil {
+ return nil, fmt.Errorf("audit_trace_store: rows: %w", err)
}
-
- statuses := make([]tracing.TraceEventStatus, 0, len(rows))
+ byTrace := map[string][]rowEvent{}
for _, r := range rows {
- statuses = append(statuses, tracing.TraceEventStatus{Source: r.Source, EventType: r.EventType})
+ byTrace[r.TraceID] = append(byTrace[r.TraceID], r)
}
- src := dominantSource(rows)
- return rootOp, src, tracing.TraceStatus(statuses), nil
+ return byTrace, nil
}
-func dominantSource(rows []rowEvent) string {
+// summarise names a trace by its root row, the first whose parent span wrote no row.
+func summarise(rows []rowEvent) (rootOp, source, status string) {
+ spans := make(map[string]bool, len(rows))
+ statuses := make([]tracing.TraceEventStatus, len(rows))
+ for i, r := range rows {
+ spans[r.SpanID] = true
+ statuses[i] = tracing.TraceEventStatus{Source: r.Source, EventType: r.EventType}
+ }
+ root := rows[0]
for _, r := range rows {
- if r.ParentSpanID == "" {
- return r.Source
+ if r.ParentSpanID == nil || !spans[*r.ParentSpanID] {
+ root = r
+ break
}
}
- return rows[0].Source
+ return root.EventType, root.Source, tracing.TraceStatus(statuses)
}
func (s *pgAuditTraceStore) GetTraceTree(ctx context.Context, traceID string) (*models.TraceNode, bool, error) {
- q := `SELECT trace_id, span_id, parent_span_id, source, event_type, entity_type, entity_id, description, created_at
- FROM (` + auditTraceSelect + `) u
- WHERE trace_id = ?
- ORDER BY created_at ASC, span_id ASC
- LIMIT ?`
- var rows []rowEvent
- if err := s.db.SelectContext(ctx, &rows, s.db.Rebind(q), traceID, TreeRowLimit+1); err != nil {
- return nil, false, fmt.Errorf("audit_trace_store: tree: %w", err)
- }
- if len(rows) == 0 {
- return nil, false, nil
+ byTrace, err := s.traceRows(ctx, `trace_id = ?`, traceID)
+ rows := byTrace[traceID]
+ if err != nil || len(rows) == 0 {
+ return nil, false, err
}
+ _, _, status := summarise(rows)
truncated := len(rows) > TreeRowLimit
- if truncated {
- rows = rows[:TreeRowLimit]
- }
-
- return buildTree(rows), truncated, nil
+ tree := buildTree(rows[:min(len(rows), TreeRowLimit)])
+ tree.Status = status
+ return tree, truncated, nil
}
-// buildTree links rows by span_id. Rows whose parent isn't in the table
-// become top-level siblings.
+// buildTree makes one node per row and hangs children off their parent span's
+// first row. Rows whose parent isn't in the table become top-level siblings.
func buildTree(rows []rowEvent) *models.TraceNode {
+ nodes := make([]*models.TraceNode, len(rows))
bySpan := make(map[string]*models.TraceNode, len(rows))
- for _, r := range rows {
- bySpan[r.SpanID] = &models.TraceNode{TraceEvent: r.toTraceEvent()}
+ for i, r := range rows {
+ nodes[i] = &models.TraceNode{TraceEvent: r.toTraceEvent()}
+ if _, ok := bySpan[r.SpanID]; !ok {
+ bySpan[r.SpanID] = nodes[i]
+ }
}
root := &models.TraceNode{}
- for _, r := range rows {
- node := bySpan[r.SpanID]
- if r.ParentSpanID != "" {
- if parent, ok := bySpan[r.ParentSpanID]; ok {
+ for i, r := range rows {
+ node := nodes[i]
+ if r.ParentSpanID != nil {
+ if parent, ok := bySpan[*r.ParentSpanID]; ok {
parent.Children = append(parent.Children, node)
continue
}
@@ -300,20 +308,13 @@ func buildTraceWhere(f TraceFilter) (string, []any) {
args = append(args, strings.ToLower(f.Q)+"%", "%"+f.Q+"%")
}
+ if f.traceIDs != nil {
+ clauses = append(clauses, "u.trace_id = ANY(?)")
+ args = append(args, f.traceIDs)
+ }
+
if len(clauses) == 0 {
return "", args
}
return "WHERE " + strings.Join(clauses, " AND "), args
}
-
-func matchesStatusFilter(status string, want []string) bool {
- if len(want) == 0 {
- return true
- }
- for _, w := range want {
- if w == status {
- return true
- }
- }
- return false
-}
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..9263c34ff
--- /dev/null
+++ b/internal/store/audit_trace_store_integration_test.go
@@ -0,0 +1,96 @@
+//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 (
+ "fmt"
+ "testing"
+ "time"
+
+ "github.com/apache/airavata-custos/pkg/models"
+)
+
+// 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, ctx := setupTestDB(t), t.Context()
+ 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 / ok.
+ for i := range 6 {
+ traceID := fmt.Sprintf("%032x", 0xa00+i)
+ eventType := []string{"ROOT_FAILED", "ROOT_STARTED"}[i%2]
+ 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)
+ }
+}
+
+// One span writing several rows yields one tree node per row, and an
+// unsourced row is stored as core.
+func TestTraceTreeKeepsEveryRowOfASpan(t *testing.T) {
+ database, ctx := setupTestDB(t), t.Context()
+ const traceID = "00000000000000000000000000000b00"
+ tx, err := database.BeginTx(ctx, nil)
+ if err != nil {
+ t.Fatalf("begin: %v", err)
+ }
+ defer tx.Rollback()
+ for i, src := range []string{"trace-span-test", ""} {
+ e := &models.AuditEvent{ID: fmt.Sprint(traceID, i), EventType: "ROOT", EventTime: time.Unix(int64(i), 0), EntityID: "e", Source: src, TraceID: traceID, SpanID: traceID[:16]}
+ if err := NewAuditEventStore(database).Create(ctx, tx, e); err != nil {
+ t.Fatalf("create: %v", err)
+ }
+ }
+ if err := tx.Commit(); err != nil {
+ t.Fatalf("commit: %v", err)
+ }
+
+ tree, _, err := NewAuditTraceStore(database).GetTraceTree(ctx, traceID)
+ if err != nil {
+ t.Fatalf("tree: %v", err)
+ }
+ if len(tree.Children) != 2 || tree.Children[0].Source != "trace-span-test" || tree.Children[1].Source != "core" {
+ t.Fatalf("tree children=%+v, want the sourced row then the core one", tree.Children)
+ }
+}
diff --git a/internal/store/audit_trace_store_test.go b/internal/store/audit_trace_store_test.go
index fb753c0dd..adf35fd71 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 {
@@ -119,24 +119,15 @@ func TestBuildTraceWhereTimeAndQ(t *testing.T) {
}
}
-func TestMatchesStatusFilter(t *testing.T) {
- if !matchesStatusFilter("ok", nil) {
- t.Error("empty filter should match")
- }
- if !matchesStatusFilter("error", []string{"ok", "error"}) {
- t.Error("match should succeed")
- }
- if matchesStatusFilter("in_progress", []string{"ok"}) {
- t.Error("filter mismatch should reject")
- }
-}
-
-func TestDominantSourcePicksRoot(t *testing.T) {
+// The root is the first row whose parent span wrote no row, so a trace
+// resumed from an unaudited span (an AMIE ingest, a publisher) still has one.
+func TestSummarisePicksRootWithUnauditedParent(t *testing.T) {
+ ingest, received := mkSpan(0x01), mkSpan(0x02)
rows := []rowEvent{
- {Source: "comanage", ParentSpanID: mkSpan(0x01)},
- {Source: "amie"}, // root: nil parent
+ {SpanID: received, ParentSpanID: &ingest, Source: "amie", EventType: "PACKET_RECEIVED"},
+ {SpanID: mkSpan(0x03), ParentSpanID: &received, Source: "comanage", EventType: "ComanageLookup"},
}
- if got := dominantSource(rows); got != "amie" {
- t.Errorf("dominant = %q, want amie", got)
+ if op, src, _ := summarise(rows); op != "PACKET_RECEIVED" || src != "amie" {
+ t.Errorf("root = %s/%s, want PACKET_RECEIVED/amie", op, src)
}
}
diff --git a/internal/store/event_delivery_store.go b/internal/store/event_delivery_store.go
index 8eba42871..f017e8683 100644
--- a/internal/store/event_delivery_store.go
+++ b/internal/store/event_delivery_store.go
@@ -30,7 +30,7 @@ import (
// Delivery columns with the event table joined in.
const eventDeliveryWithEventColumns = `d.id, d.event_id, d.subscriber, d.status, d.attempts, d.next_run_at, d.last_error, d.created_at, d.finished_at,
- e.id AS "event.id", e.event_type AS "event.event_type", e.payload AS "event.payload", e.source AS "event.source", e.trace_id AS "event.trace_id", e.created_at AS "event.created_at"`
+ e.id AS "event.id", e.event_type AS "event.event_type", e.payload AS "event.payload", e.source AS "event.source", e.trace_id AS "event.trace_id", e.span_id AS "event.span_id", e.created_at AS "event.created_at"`
type pgEventDeliveryStore struct {
db *sqlx.DB
@@ -78,9 +78,9 @@ func (s *pgEventDeliveryStore) DeleteSubscription(ctx context.Context, subscribe
func (s *pgEventDeliveryStore) CreateEvent(ctx context.Context, tx *sql.Tx, e *models.Event) error {
// Send the raw JSON as text and let Postgres cast it to jsonb.
_, err := tx.ExecContext(ctx,
- `INSERT INTO events (id, event_type, payload, source, trace_id, created_at)
- VALUES ($1, $2, $3::jsonb, $4, $5, $6)`,
- e.ID, e.EventType, string(e.Payload), e.Source, e.TraceID, e.CreatedAt)
+ `INSERT INTO events (id, event_type, payload, source, trace_id, span_id, created_at)
+ VALUES ($1, $2, $3::jsonb, $4, $5, $6, $7)`,
+ e.ID, e.EventType, string(e.Payload), e.Source, e.TraceID, e.SpanID, e.CreatedAt)
return err
}
diff --git a/internal/store/event_delivery_store_integration_test.go b/internal/store/event_delivery_store_integration_test.go
index 92e22dfcf..acf55a86c 100644
--- a/internal/store/event_delivery_store_integration_test.go
+++ b/internal/store/event_delivery_store_integration_test.go
@@ -38,6 +38,8 @@ 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",
+ SpanID: "00f067aa0ba902b7",
CreatedAt: createdAt,
}
}
diff --git a/internal/tracing/audit.go b/internal/tracing/audit.go
index 9a3fc423c..9ed7dfc89 100644
--- a/internal/tracing/audit.go
+++ b/internal/tracing/audit.go
@@ -19,16 +19,15 @@ package tracing
import "context"
-func PopulateAuditIDs(ctx context.Context, traceID, spanID, parentSpanID *string) {
- if traceID == nil || spanID == nil || parentSpanID == nil {
- return
- }
+func PopulateAuditIDs(ctx context.Context, traceID, spanID *string, parentSpanID **string) error {
if *traceID == "" && *spanID == "" {
- *traceID, *spanID = IDsFromContext(ctx)
- }
- if *parentSpanID == "" {
- if p := ParentSpanIDFromContext(ctx); p != "" {
- *parentSpanID = p
+ var err error
+ if *traceID, *spanID, err = IDsFromContext(ctx); err != nil {
+ return err
}
}
+ if *parentSpanID == nil {
+ *parentSpanID = ParentSpanIDFromContext(ctx)
+ }
+ return nil
}
diff --git a/internal/tracing/middleware.go b/internal/tracing/middleware.go
index 0cd183394..b944a84db 100644
--- a/internal/tracing/middleware.go
+++ b/internal/tracing/middleware.go
@@ -20,34 +20,34 @@ package tracing
import (
"fmt"
"net/http"
+ "strings"
"go.opentelemetry.io/otel/attribute"
"go.opentelemetry.io/otel/codes"
+ "go.opentelemetry.io/otel/propagation"
+ "go.opentelemetry.io/otel/trace"
"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 root span per request and writes X-Trace-Id. A client's
+// traceparent is only linked, so callers cannot write into another trace.
func Middleware(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
- route := r.Pattern
- if route == "" {
- route = r.URL.Path
+ var opts []trace.SpanStartOption
+ if link := trace.LinkFromContext(propagation.TraceContext{}.Extract(r.Context(), propagation.HeaderCarrier(r.Header))); link.SpanContext.IsValid() {
+ opts = append(opts, trace.WithLinks(link))
}
- name := "http." + r.Method + " " + route
-
- ctx, span := Start(r.Context(), name)
+ ctx, span := Start(r.Context(), "http."+r.Method+" "+r.URL.Path, opts...)
defer span.End()
span.SetAttributes(
attribute.String("http.method", r.Method),
- attribute.String("http.route", route),
+ attribute.String("http.route", r.URL.Path),
attribute.String("source", "http"),
)
- if tid, _ := IDsFromContext(ctx); tid != "" {
- w.Header().Set("X-Trace-Id", tid)
- }
+ w.Header().Set("X-Trace-Id", span.SpanContext().TraceID().String())
sw := &httputil.StatusRecorder{ResponseWriter: w, Status: http.StatusOK}
@@ -68,3 +68,14 @@ func Middleware(next http.Handler) http.Handler {
}
})
}
+
+// SetRoute names the request's span by the pattern, "[METHOD ]/path", that a mux matched on r.
+func SetRoute(r *http.Request) {
+ if r.Pattern == "" {
+ return
+ }
+ route := r.Pattern[strings.IndexByte(r.Pattern, ' ')+1:]
+ span := trace.SpanFromContext(r.Context())
+ span.SetName("http." + r.Method + " " + route)
+ span.SetAttributes(attribute.String("http.route", route))
+}
diff --git a/internal/tracing/middleware_test.go b/internal/tracing/middleware_test.go
index 3529f45e1..ba031e7b7 100644
--- a/internal/tracing/middleware_test.go
+++ b/internal/tracing/middleware_test.go
@@ -22,30 +22,20 @@ import (
"net/http/httptest"
"testing"
- "go.opentelemetry.io/otel"
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 {
t.Helper()
- prev := otel.GetTracerProvider()
sr := tracetest.NewSpanRecorder()
tp := sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(sr))
- otel.SetTracerProvider(tp)
- t.Cleanup(func() { otel.SetTracerProvider(prev) })
+ prev := SetProvider(tp)
+ t.Cleanup(func() { SetProvider(prev) })
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)
@@ -54,13 +44,17 @@ func TestMiddlewareProductionEmitsRootSpanAndHeader(t *testing.T) {
innerTrace string
innerSpan string
)
- mux.HandleFunc("GET /probe", func(w http.ResponseWriter, r *http.Request) {
- innerTrace, innerSpan = IDsFromContext(r.Context())
+ mux.HandleFunc("GET /probe/{id}", func(w http.ResponseWriter, r *http.Request) {
+ innerTrace, innerSpan, _ = IDsFromContext(r.Context())
w.WriteHeader(http.StatusNoContent)
})
- handler := Middleware(mux)
- req := httptest.NewRequest(http.MethodGet, "/probe", nil)
+ handler := Middleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ mux.ServeHTTP(w, r)
+ SetRoute(r)
+ }))
+ req := httptest.NewRequest(http.MethodGet, "/probe/42", nil)
+ req.Header.Set("traceparent", "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01")
rec := httptest.NewRecorder()
handler.ServeHTTP(rec, req)
@@ -75,6 +69,9 @@ func TestMiddlewareProductionEmitsRootSpanAndHeader(t *testing.T) {
if len(hdr) != 32 {
t.Fatalf("expected 32-char hex trace id, got %q (len=%d)", hdr, len(hdr))
}
+ if hdr == "4bf92f3577b34da6a3ce929d0e0e4736" {
+ t.Fatalf("client traceparent became the request's trace")
+ }
if innerTrace == "" || innerSpan == "" {
t.Fatalf("expected inner handler to see recording span; got trace=%q span=%q", innerTrace, innerSpan)
@@ -88,7 +85,10 @@ func TestMiddlewareProductionEmitsRootSpanAndHeader(t *testing.T) {
t.Fatalf("expected exactly 1 ended span, got %d", len(spans))
}
s := spans[0]
- if got, want := s.Name(), "http.GET /probe"; got != want {
+ if links := s.Links(); len(links) != 1 || links[0].SpanContext.TraceID().String() != "4bf92f3577b34da6a3ce929d0e0e4736" {
+ t.Fatalf("expected the client traceparent as the only link, got %+v", links)
+ }
+ if got, want := s.Name(), "http.GET /probe/{id}"; got != want {
t.Fatalf("span name = %q, want %q", got, want)
}
@@ -99,8 +99,8 @@ func TestMiddlewareProductionEmitsRootSpanAndHeader(t *testing.T) {
if attrs["http.method"] != "GET" {
t.Fatalf("http.method attr = %q, want GET", attrs["http.method"])
}
- if attrs["http.route"] != "/probe" {
- t.Fatalf("http.route attr = %q, want /probe", attrs["http.route"])
+ if attrs["http.route"] != "/probe/{id}" {
+ t.Fatalf("http.route attr = %q, want /probe/{id}", attrs["http.route"])
}
if attrs["http.status_code"] != "204" {
t.Fatalf("http.status_code attr = %q, want 204", attrs["http.status_code"])
@@ -110,32 +110,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/status.go b/internal/tracing/status.go
index 79f8a3264..c6f4acf63 100644
--- a/internal/tracing/status.go
+++ b/internal/tracing/status.go
@@ -18,6 +18,7 @@
package tracing
import (
+ "slices"
"strings"
"sync"
)
@@ -40,18 +41,23 @@ func EventStatus(eventType string) string {
return StatusOk
}
+type traceMarkers struct {
+ start string
+ terminals []string
+}
+
var (
- terminalMarkersMu sync.RWMutex
- terminalMarkers = map[string][]string{}
+ markersMu sync.RWMutex
+ markers = map[string]traceMarkers{}
)
-// RegisterTerminalMarkers declares the event names that close out a trace for
-// the given source. Connectors call this at boot so the core stays unaware of
-// connector-specific event names.
-func RegisterTerminalMarkers(source string, markers ...string) {
- terminalMarkersMu.Lock()
- defer terminalMarkersMu.Unlock()
- terminalMarkers[source] = append(terminalMarkers[source], markers...)
+// RegisterMarkers declares the event that opens a source's work in a trace and
+// the events that close it. Connectors call this at boot so the core stays
+// unaware of connector-specific event names.
+func RegisterMarkers(source, start string, terminals ...string) {
+ markersMu.Lock()
+ defer markersMu.Unlock()
+ markers[source] = traceMarkers{start, terminals}
}
type TraceEventStatus struct {
@@ -59,49 +65,25 @@ type TraceEventStatus struct {
EventType string
}
-func setTerminalMarkersForTest(t interface{ Cleanup(func()) }, source string, markers []string) {
- terminalMarkersMu.Lock()
- prev, hadPrev := terminalMarkers[source]
- terminalMarkers[source] = markers
- terminalMarkersMu.Unlock()
- t.Cleanup(func() {
- terminalMarkersMu.Lock()
- defer terminalMarkersMu.Unlock()
- if hadPrev {
- terminalMarkers[source] = prev
- } else {
- delete(terminalMarkers, source)
- }
- })
-}
-
-// TraceStatus is "error" if any event errored, "ok" if a registered terminal
-// marker is present, else "in_progress".
+// TraceStatus is "error" if any event errored, "in_progress" while a source has
+// written its start marker and none of its terminals, else "ok". A source with
+// no registered markers never holds a trace in progress.
func TraceStatus(events []TraceEventStatus) string {
- terminalMarkersMu.RLock()
- snapshot := make(map[string][]string, len(terminalMarkers))
- for k, v := range terminalMarkers {
- snapshot[k] = v
- }
- terminalMarkersMu.RUnlock()
-
- hasError := false
- hasTerminal := false
+ markersMu.RLock()
+ defer markersMu.RUnlock()
+ started, done := map[string]bool{}, map[string]bool{}
for _, e := range events {
if EventStatus(e.EventType) == StatusError {
- hasError = true
- }
- for _, m := range snapshot[e.Source] {
- if e.EventType == m {
- hasTerminal = true
- }
+ return StatusError
}
+ m, ok := markers[e.Source]
+ started[e.Source] = started[e.Source] || ok && e.EventType == m.start
+ done[e.Source] = done[e.Source] || slices.Contains(m.terminals, e.EventType)
}
- if hasError {
- return StatusError
- }
- if hasTerminal {
- return StatusOk
+ for s := range started {
+ if started[s] && !done[s] {
+ return StatusInProgress
+ }
}
- return StatusInProgress
+ return StatusOk
}
diff --git a/internal/tracing/status_test.go b/internal/tracing/status_test.go
index 3eb21ccc6..a03ce470a 100644
--- a/internal/tracing/status_test.go
+++ b/internal/tracing/status_test.go
@@ -36,52 +36,25 @@ func TestEventStatus(t *testing.T) {
}
}
-func TestTraceStatusOk(t *testing.T) {
- setTerminalMarkersForTest(t, "amie", []string{"TRANSACTION_COMPLETE"})
- events := []TraceEventStatus{
- {Source: "amie", EventType: "CREATE_PERSON"},
- {Source: "amie", EventType: "CREATE_ACCOUNT"},
- {Source: "amie", EventType: "TRANSACTION_COMPLETE"},
- }
- if got := TraceStatus(events); got != StatusOk {
- t.Errorf("TraceStatus(ok flow) = %q, want %q", got, StatusOk)
- }
-}
-
-func TestTraceStatusError(t *testing.T) {
- setTerminalMarkersForTest(t, "amie", []string{"TRANSACTION_COMPLETE"})
- events := []TraceEventStatus{
- {Source: "amie", EventType: "CREATE_PERSON"},
- {Source: "comanage", EventType: "ComanageProvisioningFailed"},
- {Source: "amie", EventType: "TRANSACTION_COMPLETE"},
- }
- if got := TraceStatus(events); got != StatusError {
- t.Errorf("TraceStatus(error wins over terminal) = %q, want %q", got, StatusError)
- }
-}
-
-func TestTraceStatusInProgress(t *testing.T) {
- events := []TraceEventStatus{
- {Source: "amie", EventType: "CREATE_PERSON"},
- {Source: "amie", EventType: "CREATE_ACCOUNT"},
- }
- if got := TraceStatus(events); got != StatusInProgress {
- t.Errorf("TraceStatus(no terminal) = %q, want %q", got, StatusInProgress)
- }
-}
-
-func TestTraceStatusEmpty(t *testing.T) {
- if got := TraceStatus(nil); got != StatusInProgress {
- t.Errorf("TraceStatus(nil) = %q, want %q", got, StatusInProgress)
- }
-}
-
-func TestTraceStatusComanageTerminal(t *testing.T) {
- setTerminalMarkersForTest(t, "comanage", []string{"ComanageClusterAccountAttached"})
- events := []TraceEventStatus{
- {Source: "comanage", EventType: "ComanageClusterAccountAttached"},
- }
- if got := TraceStatus(events); got != StatusOk {
- t.Errorf("TraceStatus(comanage terminal) = %q, want %q", got, StatusOk)
+func TestTraceStatus(t *testing.T) {
+ RegisterMarkers("status-test", "PACKET_RECEIVED", "TRANSACTION_COMPLETE", "REPLY_HELD")
+ cases := []struct {
+ name string
+ events []string
+ want string
+ }{
+ {"start without terminal", []string{"PACKET_RECEIVED", "CREATE_PERSON"}, StatusInProgress},
+ {"no start", []string{"CREATE_PERSON"}, StatusOk},
+ {"start and terminal", []string{"PACKET_RECEIVED", "REPLY_HELD"}, StatusOk},
+ {"error wins", []string{"PACKET_RECEIVED", "REQUEST_REJECTED", "TRANSACTION_COMPLETE"}, StatusError},
+ }
+ for _, c := range cases {
+ events := make([]TraceEventStatus, len(c.events))
+ for i, et := range c.events {
+ events[i] = TraceEventStatus{Source: "status-test", EventType: et}
+ }
+ if got := TraceStatus(events); got != c.want {
+ t.Errorf("%s: TraceStatus = %q, want %q", c.name, got, c.want)
+ }
}
}
diff --git a/internal/tracing/tracing.go b/internal/tracing/tracing.go
index 92bcdebc9..2cd05f552 100644
--- a/internal/tracing/tracing.go
+++ b/internal/tracing/tracing.go
@@ -19,7 +19,7 @@ package tracing
import (
"context"
- "crypto/rand"
+ "errors"
"fmt"
"log/slog"
"os"
@@ -31,31 +31,27 @@ 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 tracerName = "custos"
-const (
- ModeProduction Mode = iota
- ModeNoop
-)
+// provider mints the spans whose ids audit rows carry, so it is a real SDK
+// provider from the start rather than the otel global, which is a no-op until set.
+var provider trace.TracerProvider = sdktrace.NewTracerProvider()
-const tracerName = "custos"
+// SetProvider replaces the span provider, for Init and for tests that record
+// spans, and returns the previous one.
+func SetProvider(tp trace.TracerProvider) (prev trace.TracerProvider) {
+ prev, provider = provider, tp
+ return prev
+}
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
- }
-
serviceName := cfg.ServiceName
if serviceName == "" {
serviceName = "custos"
@@ -77,6 +73,7 @@ func Init(cfg InitConfig) (func(context.Context) error, error) {
// No SpanProcessor: spans live in ctx for ID propagation only.
tp := sdktrace.NewTracerProvider(sdktrace.WithResource(res))
+ SetProvider(tp)
otel.SetTracerProvider(tp)
otel.SetTextMapPropagator(propagation.TraceContext{})
@@ -96,7 +93,7 @@ var (
// Start opens a span and stamps the audit parent in ctx. bus.* spans are
// skipped so audit parents jump over the bus to the nearest business span.
func Start(ctx context.Context, name string, opts ...trace.SpanStartOption) (context.Context, trace.Span) {
- newCtx, span := otel.Tracer(tracerName).Start(ctx, name, opts...)
+ newCtx, span := provider.Tracer(tracerName).Start(ctx, name, opts...)
if strings.HasPrefix(name, "bus.") {
return newCtx, span
@@ -112,34 +109,25 @@ 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 ""
-}
-
-func FromContext(ctx context.Context) trace.Span {
- return trace.SpanFromContext(ctx)
+ return nil
}
-// 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.
-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)
- return ctx
- }
- var sid trace.SpanID
- if _, err := rand.Read(sid[:]); err != nil {
+// ContextWithSpanContext resumes a saved span, so the spans opened from here
+// on are its children and audit rows written here take it as their parent.
+// An invalid one leaves ctx unchanged and is logged since it should not happen.
+func ContextWithSpanContext(ctx context.Context, traceID, spanID string) context.Context {
+ tid, terr := trace.TraceIDFromHex(traceID)
+ sid, serr := trace.SpanIDFromHex(spanID)
+ if err := errors.Join(terr, serr); err != nil {
+ slog.WarnContext(ctx, "ignoring invalid span context", "trace_id", traceID, "span_id", spanID, "error", err)
return ctx
}
+ ctx = context.WithValue(ctx, parentSpanIDKey, sid)
return trace.ContextWithRemoteSpanContext(ctx, trace.NewSpanContext(trace.SpanContextConfig{
TraceID: tid,
SpanID: sid,
@@ -148,22 +136,13 @@ 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 "", ""
- }
- return sc.TraceID().String(), sc.SpanID().String()
-}
+// ErrNoSpan rejects a write made outside a span, since every audit row and event needs a trace.
+var ErrNoSpan = errors.New("tracing: no active span")
-// IDsBytesFromContext returns the active trace_id (16 bytes) and span_id (8 bytes),
-// or nil/nil if no recording span is on ctx.
-func IDsBytesFromContext(ctx context.Context) (traceID, spanID []byte) {
+func IDsFromContext(ctx context.Context) (traceID, spanID string, err error) {
sc := trace.SpanContextFromContext(ctx)
if !sc.IsValid() {
- return nil, nil
+ return "", "", ErrNoSpan
}
- tid := sc.TraceID()
- sid := sc.SpanID()
- return append([]byte(nil), tid[:]...), append([]byte(nil), sid[:]...)
+ return sc.TraceID().String(), sc.SpanID().String(), nil
}
diff --git a/internal/tracing/tracing_test.go b/internal/tracing/tracing_test.go
index 71ade28db..840d123d6 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)
}
@@ -62,14 +37,14 @@ func TestInitProductionMintsValidIDs(t *testing.T) {
ctx, span := Start(context.Background(), "test.production")
defer span.End()
- traceID, spanID := IDsFromContext(ctx)
+ traceID, spanID, _ := IDsFromContext(ctx)
if traceID == "" || spanID == "" {
t.Fatalf("expected non-empty IDs from production ctx, got trace=%q span=%q", traceID, spanID)
}
}
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/internal/tracing/audit.go b/internal/tracing/tracingtest/tracingtest.go
similarity index 68%
copy from internal/tracing/audit.go
copy to internal/tracing/tracingtest/tracingtest.go
index 9a3fc423c..4053f18fd 100644
--- a/internal/tracing/audit.go
+++ b/internal/tracing/tracingtest/tracingtest.go
@@ -15,20 +15,17 @@
// specific language governing permissions and limitations
// under the License.
-package tracing
+// Package tracingtest gives tests the span that every audit row and event write requires.
+package tracingtest
-import "context"
+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
- }
- }
+ "github.com/apache/airavata-custos/internal/tracing"
+)
+
+// Context returns a background context inside a fresh root span.
+func Context() context.Context {
+ ctx, _ := tracing.Start(context.Background(), "test")
+ return ctx
}
diff --git a/pkg/events/bus.go b/pkg/events/bus.go
index a24e71017..e34d83078 100644
--- a/pkg/events/bus.go
+++ b/pkg/events/bus.go
@@ -127,7 +127,7 @@ func subscribeTyped[T any](b *Bus, subscriber string, topic EventType, handler f
return handler(ctx, *v)
}
default:
- slog.Warn("event payload has unexpected type", "type", event.Type, "got", value)
+ slog.WarnContext(ctx, "event payload has unexpected type", "type", event.Type, "got", value)
}
return nil
})
@@ -147,7 +147,10 @@ func (b *Bus) Publish(ctx context.Context, tx *sql.Tx, topic EventType, payload
if err != nil {
return fmt.Errorf("encode event payload: %w", err)
}
- traceID, _ := tracing.IDsFromContext(ctx)
+ traceID, spanID, err := tracing.IDsFromContext(ctx)
+ if err != nil {
+ return fmt.Errorf("publish %s: %w", topic, err)
+ }
source := audit.SourceFromContext(ctx)
// FIXME - remove defaulting to 'core' every entry point should put its source
if source == "" {
@@ -160,6 +163,7 @@ func (b *Bus) Publish(ctx context.Context, tx *sql.Tx, topic EventType, payload
Payload: body,
Source: source,
TraceID: traceID,
+ SpanID: spanID,
CreatedAt: now,
}
@@ -212,7 +216,7 @@ func callHandler(ctx context.Context, h EventSubscriberFunc, event Event, payloa
defer func() {
if r := recover(); r != nil {
err = fmt.Errorf("subscriber panic: %v", r)
- slog.Error("event subscriber panicked", "topic", event.Type, "panic", r, "stack", string(debug.Stack()))
+ slog.ErrorContext(ctx, "event subscriber panicked", "topic", event.Type, "panic", r, "stack", string(debug.Stack()))
}
}()
return h(ctx, event, payload)
@@ -248,7 +252,7 @@ func dispatchSync(ctx context.Context, h handler, event Event, payload any) erro
err := fmt.Errorf("subscriber panic: %v", r)
span.RecordError(err)
span.SetStatus(codes.Error, "subscriber panic")
- slog.Error("event subscriber panicked",
+ slog.ErrorContext(ctx, "event subscriber panicked",
"topic", event.Type,
"subscriber", h.subscriber,
"panic", r,
diff --git a/pkg/events/bus_test.go b/pkg/events/bus_test.go
index 0c6451c00..86e309781 100644
--- a/pkg/events/bus_test.go
+++ b/pkg/events/bus_test.go
@@ -33,6 +33,7 @@ import (
"github.com/apache/airavata-custos/internal/db"
"github.com/apache/airavata-custos/internal/tracing"
+ "github.com/apache/airavata-custos/internal/tracing/tracingtest"
"github.com/apache/airavata-custos/pkg/models"
)
@@ -67,7 +68,7 @@ func publish(t *testing.T, ctx context.Context, bus *Bus, topic EventType, paylo
func auditRows(t *testing.T, database *sqlx.DB, entityID string) []models.AuditEvent {
t.Helper()
var rows []models.AuditEvent
- if err := database.Select(&rows, "SELECT event_type, details, trace_id FROM audit_events WHERE entity_id = $1 ORDER BY event_time", entityID); err != nil {
+ if err := database.Select(&rows, "SELECT event_type, details, trace_id, parent_span_id FROM audit_events WHERE entity_id = $1 ORDER BY event_time", entityID); err != nil {
t.Fatalf("list audit rows: %v", err)
}
return rows
@@ -92,12 +93,12 @@ func TestPublishIsDeliveredToSubscriber(t *testing.T) {
var gotTrace string
bus.SubscribeComputeClusterUserCreated(subscriberTest, func(ctx context.Context, cu models.ComputeClusterUser) error {
got = cu
- gotTrace, _ = tracing.IDsFromContext(ctx)
+ gotTrace, _, _ = tracing.IDsFromContext(ctx)
return nil
})
ctx, span := tracing.Start(context.Background(), "test.publish")
- wantTrace, _ := tracing.IDsFromContext(ctx)
+ wantTrace, _, _ := tracing.IDsFromContext(ctx)
publish(t, ctx, bus, ComputeClusterUserCreateEvent, clusterUser())
span.End()
bus.deliverDue(context.Background())
@@ -118,7 +119,7 @@ func TestPublishWritesNothingWithoutSubscription(t *testing.T) {
database := setupTestDB(t)
bus := newBus(t, database)
- publish(t, context.Background(), bus, UserCreateEvent, &models.User{ID: "u-1"})
+ publish(t, tracingtest.Context(), bus, UserCreateEvent, &models.User{ID: "u-1"})
if n := countEvents(t, database); n != 0 {
t.Fatalf("expected no stored events, got %d", n)
}
@@ -134,7 +135,7 @@ func TestPublishRollsBackWithCallerTransaction(t *testing.T) {
if err != nil {
t.Fatalf("begin: %v", err)
}
- if err := bus.Publish(context.Background(), tx, ComputeClusterUserCreateEvent, clusterUser()); err != nil {
+ if err := bus.Publish(tracingtest.Context(), tx, ComputeClusterUserCreateEvent, clusterUser()); err != nil {
t.Fatalf("publish: %v", err)
}
if err := tx.Rollback(); err != nil {
@@ -161,7 +162,7 @@ func TestEveryDeliveryAttemptCreatesAnAuditRow(t *testing.T) {
})
ctx, span := tracing.Start(context.Background(), "test.publish")
- wantTrace, _ := tracing.IDsFromContext(ctx)
+ wantTrace, wantParent, _ := tracing.IDsFromContext(ctx)
publish(t, ctx, bus, ComputeClusterUserCreateEvent, clusterUser())
span.End()
@@ -184,8 +185,8 @@ func TestEveryDeliveryAttemptCreatesAnAuditRow(t *testing.T) {
t.Fatalf("expected the second row to carry attempt 2, got %s", rows[1].Details)
}
for _, r := range rows {
- if r.TraceID != wantTrace {
- t.Fatalf("audit row trace %q, want the publisher's %q", r.TraceID, wantTrace)
+ if r.TraceID != wantTrace || r.ParentSpanID == nil || *r.ParentSpanID != wantParent {
+ t.Fatalf("audit row trace %q parent %v, want the publisher's %q %q", r.TraceID, r.ParentSpanID, wantTrace, wantParent)
}
}
}
@@ -206,7 +207,7 @@ func TestFailedDeliveryIsRetriedAfterBackoff(t *testing.T) {
return handler(ctx, cu)
})
- publish(t, context.Background(), bus, ComputeClusterUserCreateEvent, clusterUser())
+ publish(t, tracingtest.Context(), bus, ComputeClusterUserCreateEvent, clusterUser())
// The second pass runs before the backoff has passed, so it must skip the row.
bus.deliverDue(context.Background())
bus.deliverDue(context.Background())
@@ -234,7 +235,7 @@ func TestDeliveryFailsAtAttemptCap(t *testing.T) {
return errors.New("registry 503")
})
- publish(t, context.Background(), bus, ComputeClusterUserCreateEvent, clusterUser())
+ publish(t, tracingtest.Context(), bus, ComputeClusterUserCreateEvent, clusterUser())
// Jump the row to its last allowed attempt instead of failing it nine times.
pending := onlyDelivery(t, bus, models.EventDeliveryPending)
if err := db.TxFn(context.Background(), database, func(tx *sql.Tx) error {
@@ -257,7 +258,7 @@ func TestPermanentErrorFailsWithoutRetry(t *testing.T) {
return fmt.Errorf("%w: bad row", ErrPermanent)
})
- publish(t, context.Background(), bus, ComputeClusterUserCreateEvent, clusterUser())
+ publish(t, tracingtest.Context(), bus, ComputeClusterUserCreateEvent, clusterUser())
bus.deliverDue(context.Background())
if d := onlyDelivery(t, bus, models.EventDeliveryFailed); d.Attempts != 1 {
@@ -275,7 +276,7 @@ func TestEventPublishedBeforeSubscriberLoadsIsDelivered(t *testing.T) {
// A fresh process: the subscription is saved, the handler is not loaded yet.
bus := newBus(t, database)
- publish(t, context.Background(), bus, ComputeClusterUserCreateEvent, clusterUser())
+ publish(t, tracingtest.Context(), bus, ComputeClusterUserCreateEvent, clusterUser())
bus.deliverDue(context.Background())
if d := onlyDelivery(t, bus, models.EventDeliveryPending); d.Attempts != 0 {
t.Fatalf("expected the row to wait untouched, got %d attempts", d.Attempts)
@@ -303,7 +304,7 @@ func TestUnsubscribeStopsNewDeliveries(t *testing.T) {
if err := bus.Unsubscribe(context.Background(), subscriberTest); err != nil {
t.Fatalf("unsubscribe: %v", err)
}
- publish(t, context.Background(), bus, ComputeClusterUserCreateEvent, clusterUser())
+ publish(t, tracingtest.Context(), bus, ComputeClusterUserCreateEvent, clusterUser())
if n := countEvents(t, database); n != 0 {
t.Fatalf("expected no stored events after unsubscribe, got %d", n)
}
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..20a8c994b 100644
--- a/pkg/events/worker.go
+++ b/pkg/events/worker.go
@@ -164,14 +164,14 @@ func (b *Bus) deliver(ctx context.Context, row models.PendingDelivery) {
return
}
- // Build the handler's context. It carries the trace id and source of the
+ // Build the handler's context. It carries the span and source of the
// request that published the event, so audit rows join that trace. It is
// detached from the worker's context, so shutdown does not cut the
// handler off, and it has a timeout, so a handler that never returns
// cannot stop the worker from moving on. `WithoutCancel` has to come first,
// a timeout added straight on `ctx` (instead of `WithoutCancel`) would still be canceled by shutdown.
handlerCtx := context.WithoutCancel(ctx)
- handlerCtx = tracing.ContextWithTraceID(handlerCtx, row.Event.TraceID)
+ handlerCtx = tracing.ContextWithSpanContext(handlerCtx, row.Event.TraceID, row.Event.SpanID)
handlerCtx = audit.WithSource(handlerCtx, row.Event.Source)
handlerCtx, cancel := context.WithTimeout(handlerCtx, handlerTimeout)
defer cancel()
@@ -208,7 +208,7 @@ func (b *Bus) deliver(ctx context.Context, row models.PendingDelivery) {
// The row stays pending, so the next tick runs it again. Handlers are
// idempotent, so a repeat after a success is harmless.
if recordErr != nil {
- slog.Error("event delivery outcome not recorded", "delivery_id", row.ID, "error", recordErr)
+ slog.ErrorContext(handlerCtx, "event delivery outcome not recorded", "delivery_id", row.ID, "error", recordErr)
}
}
@@ -264,7 +264,8 @@ func retryAudit(ctx context.Context, row models.PendingDelivery, actorID string,
"previous_attempts": row.Attempts,
"last_error": row.LastError,
})
- traceID, spanID := tracing.IDsFromContext(ctx)
+ // Without a span the ids stay empty and the audit store rejects the row.
+ traceID, spanID, _ := tracing.IDsFromContext(ctx)
return &models.AuditEvent{
ID: uuid.NewString(),
EventType: auditDeliveryRetried,
@@ -289,7 +290,8 @@ func deliveryAudit(ctx context.Context, row models.PendingDelivery, attempts int
details["error"] = err.Error()
}
body, _ := json.Marshal(details)
- traceID, spanID := tracing.IDsFromContext(ctx)
+ // Without a span the ids stay empty and the audit store rejects the row.
+ traceID, spanID, _ := tracing.IDsFromContext(ctx)
return &models.AuditEvent{
ID: uuid.NewString(),
EventType: eventType,
diff --git a/pkg/identity/identity.go b/pkg/identity/identity.go
index b270c9aec..c0adcf946 100644
--- a/pkg/identity/identity.go
+++ b/pkg/identity/identity.go
@@ -43,6 +43,10 @@ var (
ErrUserAlreadyOIDCLinked = fmt.Errorf("%w: user already has an OIDC binding", ErrNotLinked)
)
+type TokenVerifier interface {
+ Verify(ctx context.Context, rawToken string) (*Claims, error)
+}
+
// UserResolver maps verified OIDC claims to a Caller and the effective
// privilege set.
type UserResolver interface {
diff --git a/pkg/identity/middleware.go b/pkg/identity/middleware.go
index 9e984aa9b..d8778d7c9 100644
--- a/pkg/identity/middleware.go
+++ b/pkg/identity/middleware.go
@@ -30,7 +30,7 @@ import (
// Middleware verifies the bearer token once, resolves the caller, and
// attaches Caller + privilege set to the request context. Public paths
// bypass verification entirely.
-func Middleware(verifier *JWTVerifier, resolver UserResolver, publicPaths []string, next http.Handler) http.Handler {
+func Middleware(verifier TokenVerifier, resolver UserResolver, publicPaths []string, next http.Handler) http.Handler {
publicSet := make(map[string]struct{}, len(publicPaths))
for _, path := range publicPaths {
publicSet[path] = struct{}{}
@@ -56,7 +56,7 @@ func Middleware(verifier *JWTVerifier, resolver UserResolver, publicPaths []stri
return
}
if err != nil {
- slog.Error("identity resolution failed", "sub", claims.Sub, "email", claims.Email, "error", err.Error())
+ slog.ErrorContext(req.Context(), "identity resolution failed", "sub", claims.Sub, "email", claims.Email, "error", err.Error())
writeJSONError(w, http.StatusServiceUnavailable, "auth_lookup_failed", "Identity resolution failed")
return
}
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/event_delivery.go b/pkg/models/event_delivery.go
index 78a236734..e3c3a7739 100644
--- a/pkg/models/event_delivery.go
+++ b/pkg/models/event_delivery.go
@@ -36,6 +36,7 @@ type Event struct {
Payload json.RawMessage `json:"payload" db:"payload" swaggertype:"object"`
Source string `json:"source" db:"source"` // Subsystem that published the event.
TraceID string `json:"trace_id" db:"trace_id"`
+ SpanID string `json:"span_id" db:"span_id"` // Publisher's span, the parent of its deliveries.
CreatedAt time.Time `json:"created_at" db:"created_at"`
}
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..259751498 100644
--- a/pkg/service/audit_event.go
+++ b/pkg/service/audit_event.go
@@ -21,17 +21,10 @@ 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"
)
-// SourceCore is the default value stamped on audit rows when nothing
-// else has tagged the ctx direct service writes, admin endpoints, etc.
-const SourceCore = "core"
-
// CreateAuditEvent records a new audit event. EventTime defaults to the
// server's current UTC time when unset. Audit events are append-only — there
// is no update operation.
@@ -53,21 +46,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
- } else {
- e.Source = SourceCore
- }
- }
-
if err := s.inTx(ctx, func(tx *sql.Tx) error {
return s.auditEvents.Create(ctx, tx, e)
}); err != nil {
diff --git a/pkg/service/audit_event_test.go b/pkg/service/audit_event_test.go
index feccf9819..673ad724d 100644
--- a/pkg/service/audit_event_test.go
+++ b/pkg/service/audit_event_test.go
@@ -21,19 +21,11 @@ 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()
@@ -51,21 +43,7 @@ func TestAuditEventTraceIDs_PopulatedFromActiveSpan(t *testing.T) {
}
}
-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()
diff --git a/pkg/service/identity_resolver.go b/pkg/service/identity_resolver.go
index 5b80799bc..85a86fd4e 100644
--- a/pkg/service/identity_resolver.go
+++ b/pkg/service/identity_resolver.go
@@ -135,7 +135,7 @@ func (s *Service) ResolveCaller(ctx context.Context, claims *identity.Claims) (*
// linkBySub creates the first OIDC binding for a PENDING user matched by verified email.
func (s *Service) linkBySub(ctx context.Context, claims *identity.Claims) (*models.User, error) {
- slog.Debug("linkBySub: evaluating claims",
+ slog.DebugContext(ctx, "linkBySub: evaluating claims",
"sub", claims.Sub,
"email", claims.Email,
"email_verified", claims.EmailVerified,
@@ -164,7 +164,7 @@ func (s *Service) linkBySub(ctx context.Context, claims *identity.Claims) (*mode
// A non-MERGED user without an OIDC binding is an inconsistent state;
// merged users legitimately lose theirs to the surviving user.
if user.Status != models.UserMerged {
- slog.Error("inconsistent user state: reached non-PENDING lifecycle without an OIDC binding",
+ slog.ErrorContext(ctx, "inconsistent user state: reached non-PENDING lifecycle without an OIDC binding",
"user_id", user.ID,
"email", claims.Email,
"sub", claims.Sub,
@@ -200,7 +200,7 @@ func (s *Service) linkBySub(ctx context.Context, claims *identity.Claims) (*mode
}); err != nil {
return nil, err
}
- slog.Info("identity linked via email fallback", "user_id", user.ID, "email", claims.Email)
+ slog.InfoContext(ctx, "identity linked via email fallback", "user_id", user.ID, "email", claims.Email)
user.Status = models.UserActive
return user, nil
diff --git a/pkg/service/integration_common_test.go b/pkg/service/integration_common_test.go
index 0cb9b7a23..d62c84f80 100644
--- a/pkg/service/integration_common_test.go
+++ b/pkg/service/integration_common_test.go
@@ -30,6 +30,7 @@ import (
"github.com/jmoiron/sqlx"
"github.com/apache/airavata-custos/internal/db"
+ "github.com/apache/airavata-custos/internal/tracing/tracingtest"
"github.com/apache/airavata-custos/pkg/events"
"github.com/apache/airavata-custos/pkg/models"
)
@@ -160,4 +161,4 @@ func countAuditEventsOfType(t *testing.T, database *sqlx.DB, eventType, entityID
return n
}
-func ctx() context.Context { return context.Background() }
+func ctx() context.Context { return tracingtest.Context() }
diff --git a/pkg/service/user_onboard.go b/pkg/service/user_onboard.go
index 4733e773a..797d59748 100644
--- a/pkg/service/user_onboard.go
+++ b/pkg/service/user_onboard.go
@@ -274,7 +274,7 @@ func (s *Service) OnboardUser(ctx context.Context, in OnboardUserInput) (*models
}
if in.ClusterAdmin {
- slog.Info("cluster admin account created", "user_id", created.ID, "cluster_id", clusterUser.ComputeClusterID, "actor_id", in.OnboardedBy)
+ slog.InfoContext(ctx, "cluster admin account created", "user_id", created.ID, "cluster_id", clusterUser.ComputeClusterID, "actor_id", in.OnboardedBy)
}
return created, nil
diff --git a/pkg/service/user_role.go b/pkg/service/user_role.go
index 06f5f3189..32d529cca 100644
--- a/pkg/service/user_role.go
+++ b/pkg/service/user_role.go
@@ -180,7 +180,7 @@ func (s *Service) BootstrapSuperAdmin(ctx context.Context, email, source string)
return fmt.Errorf("lookup existing assignment: %w", err)
}
if existing != nil {
- slog.Info("bootstrap: super_admin already granted to user, skipping", "email", email)
+ slog.InfoContext(ctx, "bootstrap: super_admin already granted to user, skipping", "email", email)
return nil
}
// If another user already holds super_admin, install is past bootstrap.
@@ -189,7 +189,7 @@ func (s *Service) BootstrapSuperAdmin(ctx context.Context, email, source string)
return fmt.Errorf("count super_admin holders: %w", err)
}
if len(holders) > 0 {
- slog.Info("bootstrap: super_admin already held by another user, skipping", "email", email)
+ slog.InfoContext(ctx, "bootstrap: super_admin already held by another user, skipping", "email", email)
return nil
}
assignment := &models.UserRole{
@@ -209,7 +209,7 @@ func (s *Service) BootstrapSuperAdmin(ctx context.Context, email, source string)
}); err != nil {
return fmt.Errorf("audit bootstrap assignment: %w", err)
}
- slog.Info("bootstrap: super_admin granted", "user_id", user.ID, "email", email, "source", source)
+ slog.InfoContext(ctx, "bootstrap: super_admin granted", "user_id", user.ID, "email", email, "source", source)
return nil
})
}
@@ -241,7 +241,7 @@ func (s *Service) createBootstrapUser(ctx context.Context, email, source string)
if err != nil {
return nil, err
}
- slog.Info("bootstrap: created PENDING super_admin user", "user_id", created.ID, "email", email)
+ slog.InfoContext(ctx, "bootstrap: created PENDING super_admin user", "user_id", created.ID, "email", email)
return created, nil
}