chore(tracing-utils): add tracing utils for custom instrument

[NathanFlurry]

No description provided.

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Preview	Comments	Updated (UTC)
rivetkit-serverless	Ready	Preview	Comment	Oct 31, 2025 1:11am

3 Skipped Deployments

Project	Deployment	Preview	Comments	Updated (UTC)
rivet-cloud	Ignored			Oct 31, 2025 1:11am
rivet-inspector	Ignored	Preview		Oct 31, 2025 1:11am
rivet-site	Ignored	Preview		Oct 31, 2025 1:11am

[NathanFlurry]

• fix(epoxy): handle incomplete ranges in epoxy recover key #3296
• fix(epoxy): fix replica recovery to skip old ballots #3295
• feat(epoxy): add ability to get & set coordinator state #3294
• chore(tracing-utils): add tracing utils for custom instrument #3293 👈 (View in Graphite)
• chore(guard): add metadata to rate limit error #3292
• main

---

How to use the Graphite Merge Queue

Add the label merge-queue to this PR to add it to the merge queue.

You must have a Graphite account in order to use the merge queue. Sign up using this link.

An organization admin has enabled the Graphite Merge Queue in this repository.

Please do not merge from GitHub as this will restart CI on PRs being processed by the merge queue.

This stack of pull requests is managed by Graphite. Learn more about stacking.

[pkg-pr-new]

More templates

• example-ai-agent
• example-background-jobs
• example-better-auth-external-db
• example-bots
• chat-room
• example-cloudflare-workers
• example-cloudflare-workers-hono
• example-counter
• example-counter-serverless
• example-crdt
• example-cursors
• example-cursors-raw-websocket
• example-database
• example-deno
• example-sqlite
• example-elysia
• example-express
• example-freestyle
• example-game
• example-hono
• example-hono-bun
• example-hono-react
• example-kitchen-sink
• example-next-js
• example-rate
• example-raw-fetch-handler
• example-raw-websocket-handler
• example-raw-websocket-handler-proxy
• example-react
• example-smoke-test
• starter
• example-stream
• example-sync
• example-tenant
• example-trpc
• example-workflows

@rivetkit/actor

pnpm add https://pkg.pr.new/rivet-dev/rivet/@rivetkit/actor@3293

@rivetkit/cloudflare-workers

pnpm add https://pkg.pr.new/rivet-dev/rivet/@rivetkit/cloudflare-workers@3293

@rivetkit/core

pnpm add https://pkg.pr.new/rivet-dev/rivet/@rivetkit/core@3293

@rivetkit/db

pnpm add https://pkg.pr.new/rivet-dev/rivet/@rivetkit/db@3293

@rivetkit/framework-base

pnpm add https://pkg.pr.new/rivet-dev/rivet/@rivetkit/framework-base@3293

@rivetkit/next-js

pnpm add https://pkg.pr.new/rivet-dev/rivet/@rivetkit/next-js@3293

@rivetkit/react

pnpm add https://pkg.pr.new/rivet-dev/rivet/@rivetkit/react@3293

rivetkit

pnpm add https://pkg.pr.new/rivet-dev/rivet/rivetkit@3293

@rivetkit/sql-loader

pnpm add https://pkg.pr.new/rivet-dev/rivet/@rivetkit/sql-loader@3293

@rivetkit/engine-runner

pnpm add https://pkg.pr.new/rivet-dev/rivet/@rivetkit/engine-runner@3293

@rivetkit/engine-runner-protocol

pnpm add https://pkg.pr.new/rivet-dev/rivet/@rivetkit/engine-runner-protocol@3293

commit: f6e4232

[claude]

Code Review

Summary

This PR extracts tracing utilities from rivet-util into a new dedicated package rivet-tracing-utils. The refactoring improves modularity and allows other packages like universaldb to use these utilities without circular dependencies.

Positive Aspects

Good refactoring: Extracting the tracing utilities into a separate package is a clean architectural decision that reduces coupling.

Consistent usage: The custom_instrument pattern is already extensively used throughout the codebase with 30+ usages in gasoline alone.

Proper re-export: The rivet-util/src/future.rs correctly re-exports the utilities maintaining backward compatibility.

Issues and Concerns

1. Unsafe Code Requires Careful Review

Location: engine/packages/tracing-utils/src/lib.rs:51-52

The use of get_unchecked_mut and Pin::new_unchecked relies on invariants that should be documented.

Recommendation: Add safety comments explaining why these operations are safe. This is likely copied from existing patterns but explicit safety documentation is important for maintainability.

2. Metadata Cloning May Be Unnecessary

Location: engine/packages/tracing-utils/src/lib.rs:54

The Metadata type is typically &static so cloning may be unnecessary. Consider using a reference instead to avoid an unnecessary allocation.

3. Potential Metric Recording Overhead

Location: engine/packages/tracing-utils/src/lib.rs:57-68

The metric is only recorded on Poll::Ready which is good. However every polled future records a metric which could create high cardinality with the location and name attributes.

Questions:

• Have you tested the performance impact of this instrumentation?
• Is there any aggregation or sampling strategy for high-frequency futures?
• Consider if this should be gated behind a feature flag or tracing level check.

4. Missing Documentation

The new package lacks:

• Package-level documentation explaining when to use CustomInstrumentExt vs standard Instrument
• Documentation on the metrics being recorded
• Examples of proper usage

5. Meter Name Inconsistency

Location: engine/packages/tracing-utils/src/lib.rs:83

The meter is named rivet-util-core but this is now in the rivet-tracing-utils package. Consider updating to meter rivet-tracing-utils to ensure metrics are properly attributed.

6. No Tests

The new package has no tests. While the code is extracted from existing working code having unit tests would help ensure the custom instrumentation works correctly and the unsafe code maintains Pin guarantees.

Minor Observations

• The workspace dependency setup in Cargo.toml is correct
• The universaldb usage at line 35 properly wraps the transaction retry with instrumentation
• The new package follows the repository naming conventions

Recommendations Summary

1. High Priority: Add safety comments for unsafe code blocks
2. Medium Priority: Remove unnecessary clone on metadata
3. Medium Priority: Add package and API documentation
4. Low Priority: Consider performance implications and possibly add feature flags
5. Low Priority: Fix meter name to match package name
6. Low Priority: Add basic unit tests

Overall this is a solid refactoring that improves the codebase structure. The main concerns are around documentation and ensuring the unsafe code is well-justified.

[claude]

PR Review: Add tracing utils for custom instrument

Summary

This PR extracts tracing instrumentation utilities from rivet-util into a new dedicated package rivet-tracing-utils. This is a good refactoring that improves modularity and reduces circular dependencies.

Code Quality & Best Practices ✅

Positives:

• Good separation of concerns: Moving tracing utilities to a dedicated package follows the single responsibility principle
• Clean re-export pattern: The rivet-util module now cleanly re-exports from rivet-tracing-utils, maintaining backward compatibility
• Proper workspace integration: Correctly added to workspace dependencies in root Cargo.toml
• Follows existing patterns: Uses workspace dependencies correctly with .workspace = true

Minor Suggestions:

1. Meter name consistency (tracing-utils/src/lib.rs:83):

   static ref METER: Meter = meter("rivet-util-core");

   The meter is named "rivet-util-core" but this is now in rivet-tracing-utils. Consider renaming to "rivet-tracing-utils" for consistency.

2. Missing module documentation: The new package lacks top-level documentation. Consider adding:

   //! Tracing utilities for custom instrumentation of futures.
   //!
   //! This package provides utilities for instrumenting async futures with
   //! custom tracing spans and automatic duration metrics collection.
   

Potential Bugs or Issues ⚠️

Critical - Unsafe Code (tracing-utils/src/lib.rs:51-52):

let this = unsafe { self.get_unchecked_mut() };
let inner = unsafe { Pin::new_unchecked(&mut this.inner) };

While this code appears to be correct (following standard Future implementation patterns), the safety invariants should be documented:

• get_unchecked_mut: Safe because we own the Pin<&mut Self> and won't move the data
• Pin::new_unchecked: Safe because Instrumented properly implements Unpin or maintains pinning guarantees

Recommendation: Add safety comments explaining why these unsafe operations are sound.

Metadata Cloning (tracing-utils/src/lib.rs:54):

let metadata = inner.span().metadata().clone();

The metadata is cloned on every poll, which could be inefficient for frequently-polled futures. However, Metadata is likely a cheap Arc-based type, so this is probably acceptable. Consider adding a comment noting this tradeoff.

Performance Considerations 📊

1. Metrics overhead: Recording metrics on every future completion adds overhead. This is acceptable for instrumented futures but users should be aware this is meant for strategic instrumentation, not wrapping every tiny future.

2. String allocation (tracing-utils/src/lib.rs:63):

   KeyValue::new("location", format!("{file}:{line}"))

   Allocates a new string for each completed future. Could be optimized by caching location strings, but this is likely premature optimization.

3. Polling overhead: The custom poll implementation adds minimal overhead (metadata clone + conditional metric recording only on Poll::Ready).

Security Concerns 🔒

No security issues identified. The code:

• Uses standard tracing and metrics libraries
• Doesn't handle sensitive data
• Properly scopes unsafe operations

Test Coverage 🧪

Missing: No tests were added for the new package. Consider adding:

1. Unit tests for CustomInstrumentExt:

   #[cfg(test)]
   mod tests {
       use super::*;
       
       #[tokio::test]
       async fn test_custom_instrument_records_metrics() {
           // Test that metrics are recorded on completion
       }
       
       #[tokio::test]
       async fn test_select_all_or_wait_with_empty_iter() {
           // Test the pending behavior
       }
   }

2. Integration test verifying the universaldb usage actually emits the expected metrics

Additional Observations

1. Good usage example in universaldb/src/database.rs:35: The custom_instrument usage demonstrates the intended pattern well.

2. Version bump: All package versions changed from 25.8.2 to 2.0.22-rc.1, suggesting a broader versioning change - this appears intentional.

Verdict

Approve with suggestions ✅

This is a solid refactoring that improves code organization. The main recommendations are:

1. Document unsafe code blocks with safety comments
2. Rename the meter for consistency
3. Add module-level documentation
4. Consider adding tests (though not blocking)

The code is production-ready but would benefit from these improvements for long-term maintainability.

[claude]

Code Review

This PR refactors tracing instrumentation utilities into a new dedicated package. Overall, the refactoring is well-structured, but I've identified several issues that should be addressed.

🔴 Critical Issues

1. Unsafe Code Without Safety Documentation

// engine/packages/tracing-utils/src/lib.rs:51-52
let this = unsafe { self.get_unchecked_mut() };
let inner = unsafe { Pin::new_unchecked(&mut this.inner) };

Issue: The use of unsafe blocks in the poll implementation lacks safety documentation explaining why it's sound.

Recommendation: Add SAFETY comments above each unsafe block explaining:

• Why get_unchecked_mut is safe (e.g., "we never move out of self")
• Why Pin::new_unchecked is safe (e.g., "the inner field is never unpinned")

This is especially important for Pin-related unsafe code, which has subtle correctness requirements.

---

2. Incorrect Meter Name

// engine/packages/tracing-utils/src/lib.rs:83
static ref METER: Meter = meter("rivet-util-core");

Issue: The meter name is "rivet-util-core" but this code is now in the rivet-tracing-utils package. This creates confusion for metric attribution and violates the principle of accurate naming.

Recommendation: Change to:

static ref METER: Meter = meter("rivet-tracing-utils");

---

⚠️ Moderate Issues

3. Metadata Cloning Overhead

// engine/packages/tracing-utils/src/lib.rs:54
let metadata = inner.span().metadata().clone();

Issue: Cloning metadata on every poll call adds unnecessary overhead. Metadata is static and can be accessed when Poll::Ready is returned.

Recommendation: Move the metadata access inside the Poll::Ready branch:

match inner.poll(cx) {
    Poll::Ready(val) => {
        if let Some(metadata) = inner.span().metadata() {
            if let (Some(file), Some(line)) = (metadata.file(), metadata.line()) {
                // record metrics
            }
        }
        Poll::Ready(val)
    }
    Poll::Pending => Poll::Pending,
}

Note: metadata() returns Option<&'static Metadata>, so cloning is not needed at all.

---

4. Missing Documentation on Public APIs

// engine/packages/tracing-utils/src/lib.rs:31-38
pub trait CustomInstrumentExt: Sized {
    fn custom_instrument(self, span: tracing::Span) -> CustomInstrumented<Self> { ... }
}

pub struct CustomInstrumented<T> { ... }

Issue: Public trait and struct lack documentation explaining their purpose, usage, and how they differ from standard Instrument.

Recommendation: Add comprehensive doc comments:

• What does CustomInstrumentExt do?
• How does it differ from tracing::Instrument?
• What metrics does it collect?
• When should developers use this vs standard instrumentation?

---

ℹ️ Minor Issues & Suggestions

5. Test Coverage

Observation: No tests are included for the new rivet-tracing-utils package.

Recommendation: Add unit tests for:

• select_all_or_wait with empty and non-empty iterators
• CustomInstrumented future behavior
• Metric recording (using a test meter/observer)

---

6. Redundant Import

// engine/packages/util/src/future.rs:1
pub use rivet_tracing_utils::*;

Observation: The entire module now just re-exports the tracing-utils crate. Consider whether rivet_util::future is still the right place for these utilities, or if callers should directly import from rivet_tracing_utils.

Recommendation: Document why this re-export exists (backwards compatibility?) or consider deprecating this module path in favor of direct imports.

---

7. Database Instrumentation Context

// engine/packages/universaldb/src/database.rs:35
.custom_instrument(tracing::info_span!("run_attempt"))

Observation: The span name "run_attempt" provides limited context. Consider whether additional span attributes would be helpful (e.g., retry attempt number).

Recommendation: This is minor, but consider whether the span could include more diagnostic information for debugging retry behavior.

---

✅ Positive Observations

1. Good refactoring structure: Moving shared tracing utilities to a dedicated package improves modularity
2. Proper workspace integration: Dependencies are correctly configured using workspace = true
3. Consistent with project patterns: Follows the established package structure in engine/packages/
4. Performance-aware design: Using histogram metrics with micro buckets for precise duration tracking

---

Summary

Blocking issues:

• Add SAFETY documentation for unsafe code
• Fix incorrect meter name

Recommended improvements:

• Remove unnecessary metadata cloning
• Add public API documentation
• Add test coverage

The refactoring itself is sound, but the implementation needs these fixes before merging.