1use std::fmt;
49use std::sync::atomic::{AtomicU64, Ordering};
50use std::sync::Mutex;
51use std::time::{SystemTime, UNIX_EPOCH};
52
53use super::resource::{
54 Access, AllocTag, BlockId, DeviceBlock, DeviceMemoryResource, Generation,
55 ResourceBudgetSnapshot, ResourceError, ResourceResult, StreamId,
56};
57
58#[derive(Clone, Copy, Debug, Eq, PartialEq)]
62pub enum LogAction {
63 Allocate,
64 Deallocate,
65 ReapPending,
66}
67
68impl fmt::Display for LogAction {
69 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
70 match self {
71 LogAction::Allocate => f.write_str("allocate"),
72 LogAction::Deallocate => f.write_str("deallocate"),
73 LogAction::ReapPending => f.write_str("reap_pending"),
74 }
75 }
76}
77
78#[derive(Clone, Debug, Eq, PartialEq)]
83pub enum LogResult {
84 Ok,
85 Err { kind: &'static str, message: String },
86}
87
88impl LogResult {
89 pub fn from_result<T>(r: &ResourceResult<T>) -> Self {
92 match r {
93 Ok(_) => LogResult::Ok,
94 Err(e) => LogResult::Err {
95 kind: classify_error(e),
96 message: truncate(format!("{}", e), 256),
97 },
98 }
99 }
100}
101
102fn classify_error(e: &ResourceError) -> &'static str {
103 match e {
104 ResourceError::AllocationRetained(_) => "AllocationRetained",
105 ResourceError::OutOfBudget { .. } => "OutOfBudget",
106 ResourceError::Driver(_) => "Driver",
107 ResourceError::StreamMisuse(_) => "StreamMisuse",
108 ResourceError::UseAfterFree { .. } => "UseAfterFree",
109 ResourceError::OutOfBounds { .. } => "OutOfBounds",
110 }
111}
112
113fn truncate(mut s: String, cap: usize) -> String {
114 if s.len() > cap {
115 let mut end = cap;
122 while end > 0 && !s.is_char_boundary(end) {
123 end -= 1;
124 }
125 s.truncate(end);
126 s.push('…');
127 }
128 s
129}
130
131#[derive(Clone, Debug)]
134pub struct LogRecord {
135 pub action: LogAction,
136 pub device_ordinal: u32,
137 pub stream_id: Option<StreamId>,
138 pub ptr: Option<u64>,
139 pub bytes: Option<usize>,
140 pub tag: Option<AllocTag>,
141 pub generation: Option<Generation>,
142 pub thread_id: u64,
143 pub order_counter: u64,
144 pub timestamp_nanos: u128,
145 pub result: LogResult,
146}
147
148static ORDER_COUNTER: AtomicU64 = AtomicU64::new(1);
153
154fn next_order_counter() -> u64 {
155 ORDER_COUNTER.fetch_add(1, Ordering::Relaxed)
156}
157
158fn now_nanos() -> u128 {
159 SystemTime::now()
160 .duration_since(UNIX_EPOCH)
161 .map(|d| d.as_nanos())
162 .unwrap_or(0)
163}
164
165fn current_thread_id_u64() -> u64 {
166 let s = format!("{:?}", std::thread::current().id());
170 let mut h: u64 = 0xcbf2_9ce4_8422_2325;
171 for b in s.as_bytes() {
172 h ^= *b as u64;
173 h = h.wrapping_mul(0x100_0000_01b3);
174 }
175 h
176}
177
178#[derive(Debug)]
184pub enum SinkError {
185 Refused(String),
187 Io(String),
189}
190
191impl fmt::Display for SinkError {
192 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
193 match self {
194 SinkError::Refused(m) => write!(f, "log sink refused record: {}", m),
195 SinkError::Io(m) => write!(f, "log sink io error: {}", m),
196 }
197 }
198}
199
200impl std::error::Error for SinkError {}
201
202pub trait LoggingSink: Send + Sync {
206 fn emit(&self, record: LogRecord) -> Result<(), SinkError>;
209}
210
211pub struct InMemorySink {
215 records: Mutex<Vec<LogRecord>>,
216}
217
218impl InMemorySink {
219 pub fn new() -> Self {
220 Self {
221 records: Mutex::new(Vec::new()),
222 }
223 }
224
225 pub fn snapshot(&self) -> Vec<LogRecord> {
227 self.records.lock().expect("InMemorySink poisoned").clone()
228 }
229
230 pub fn clear(&self) {
233 self.records.lock().expect("InMemorySink poisoned").clear();
234 }
235
236 pub fn len(&self) -> usize {
238 self.records.lock().expect("InMemorySink poisoned").len()
239 }
240
241 pub fn is_empty(&self) -> bool {
242 self.len() == 0
243 }
244}
245
246impl Default for InMemorySink {
247 fn default() -> Self {
248 Self::new()
249 }
250}
251
252impl LoggingSink for InMemorySink {
253 fn emit(&self, record: LogRecord) -> Result<(), SinkError> {
254 self.records
255 .lock()
256 .expect("InMemorySink poisoned")
257 .push(record);
258 Ok(())
259 }
260}
261
262pub struct NullSink;
278
279impl NullSink {
280 pub fn new() -> Self {
281 Self
282 }
283}
284
285impl Default for NullSink {
286 fn default() -> Self {
287 Self
288 }
289}
290
291impl LoggingSink for NullSink {
292 fn emit(&self, _record: LogRecord) -> Result<(), SinkError> {
293 Ok(())
294 }
295}
296
297pub struct LoggingResource {
299 inner: std::sync::Arc<dyn DeviceMemoryResource + Send + Sync>,
300 sink: std::sync::Arc<dyn LoggingSink>,
301 dropped_records: AtomicU64,
305}
306
307impl LoggingResource {
308 pub fn new(
311 inner: Box<dyn DeviceMemoryResource + Send + Sync>,
312 sink: std::sync::Arc<dyn LoggingSink>,
313 ) -> Self {
314 Self {
315 inner: std::sync::Arc::from(inner),
316 sink,
317 dropped_records: AtomicU64::new(0),
318 }
319 }
320
321 pub fn dropped_records(&self) -> u64 {
324 self.dropped_records.load(Ordering::Relaxed)
325 }
326
327 fn emit(&self, record: LogRecord) {
328 if self.sink.emit(record).is_err() {
329 self.dropped_records.fetch_add(1, Ordering::Relaxed);
330 }
331 }
332
333 fn record_allocation(
334 &self,
335 bytes: usize,
336 stream: StreamId,
337 tag: AllocTag,
338 result: &ResourceResult<crate::memory::ResourceBlockRetirement>,
339 ) {
340 let (ptr, generation, recorded_bytes) = match &result {
341 Ok(owner) => {
342 let block = owner.block();
343 (Some(block.ptr), Some(block.generation), Some(block.bytes))
344 }
345 Err(_) => (None, None, Some(bytes)),
346 };
347 self.emit(LogRecord {
348 action: LogAction::Allocate,
349 device_ordinal: self.inner.device_ordinal(),
350 stream_id: Some(stream),
351 ptr,
352 bytes: recorded_bytes,
353 tag: Some(tag),
354 generation,
355 thread_id: current_thread_id_u64(),
356 order_counter: next_order_counter(),
357 timestamp_nanos: now_nanos(),
358 result: LogResult::from_result(result),
359 });
360 }
361}
362
363impl DeviceMemoryResource for LoggingResource {
364 fn materialize(
365 &self,
366 request: super::resource::AllocationRequest,
367 ) -> ResourceResult<DeviceBlock> {
368 let (bytes, stream, tag) = (request.bytes, request.stream, request.tag);
369 let reclamation = request.reclamation();
370 let result = self.inner.materialize(request).map(|block| {
371 crate::memory::ResourceBlockRetirement::new(
372 block,
373 std::sync::Arc::clone(&self.inner),
374 reclamation,
375 )
376 });
377 self.record_allocation(bytes, stream, tag, &result);
378 result.map(crate::memory::ResourceBlockRetirement::into_block)
379 }
380
381 fn allocation_accounting(&self) -> std::sync::Arc<super::resource::AllocationAccounting> {
382 self.inner.allocation_accounting()
383 }
384
385 fn deallocate(&self, block: DeviceBlock) -> ResourceResult<()> {
386 let ptr = block.ptr;
390 let bytes = block.bytes;
391 let tag = block.tag;
392 let gen = block.generation;
393 let stream = block.alloc_stream;
394 let dev = block.device_ordinal;
395
396 let result = self.inner.deallocate(block);
397 self.emit(LogRecord {
398 action: LogAction::Deallocate,
399 device_ordinal: dev,
400 stream_id: Some(stream),
401 ptr: Some(ptr),
402 bytes: Some(bytes),
403 tag: Some(tag),
404 generation: Some(gen),
405 thread_id: current_thread_id_u64(),
406 order_counter: next_order_counter(),
407 timestamp_nanos: now_nanos(),
408 result: LogResult::from_result(&result),
409 });
410 result
411 }
412
413 fn device_ordinal(&self) -> u32 {
414 self.inner.device_ordinal()
415 }
416
417 fn bytes_outstanding(&self) -> usize {
418 self.inner.bytes_outstanding()
419 }
420
421 fn budget_snapshot(&self) -> Option<ResourceBudgetSnapshot> {
422 self.inner.budget_snapshot()
423 }
424
425 fn reap_pending(&self) -> ResourceResult<()> {
426 let result = self.inner.reap_pending();
427 self.emit(LogRecord {
428 action: LogAction::ReapPending,
429 device_ordinal: self.inner.device_ordinal(),
430 stream_id: None,
431 ptr: None,
432 bytes: None,
433 tag: None,
434 generation: None,
435 thread_id: current_thread_id_u64(),
436 order_counter: next_order_counter(),
437 timestamp_nanos: now_nanos(),
438 result: LogResult::from_result(&result),
439 });
440 result
441 }
442
443 fn record_block_use(&self, block: &DeviceBlock, use_stream: StreamId) -> ResourceResult<()> {
444 self.inner.record_block_use(block, use_stream)
453 }
454
455 fn supports_block_use_tracking(&self) -> bool {
456 self.inner.supports_block_use_tracking()
457 }
458
459 fn access_dependencies(
460 &self,
461 block: BlockId,
462 bytes: usize,
463 ) -> ResourceResult<Option<std::sync::Arc<super::resource::DeviceAccessDependencies>>> {
464 self.inner.access_dependencies(block, bytes)
465 }
466
467 fn prepare_block_use(
468 &self,
469 block: BlockId,
470 use_stream: StreamId,
471 access: Access,
472 ) -> ResourceResult<()> {
473 self.inner.prepare_block_use(block, use_stream, access)
476 }
477
478 fn finish_block_use(
479 &self,
480 block: BlockId,
481 use_stream: StreamId,
482 access: Access,
483 ) -> ResourceResult<()> {
484 self.inner.finish_block_use(block, use_stream, access)
486 }
487}
488
489#[cfg(test)]
490mod tests {
491 use super::super::direct::DirectCudaResource;
492 use super::super::resource::BlockState;
493 use super::*;
494 use std::sync::Arc;
495
496 #[test]
497 fn host_logging_panic_retires_the_unpublished_allocation() {
498 struct PanicSink;
499 impl LoggingSink for PanicSink {
500 fn emit(&self, _: LogRecord) -> Result<(), SinkError> {
501 panic!("telemetry publication interrupted");
502 }
503 }
504 let _serial = crate::cuda_graph::capture_lifecycle_test_guard();
505 let inner = super::super::budget::tests::DeferredHostResource {
506 release_on_detach: true,
507 ..Default::default()
508 };
509 let ledger = Arc::clone(&inner.accounting);
510 let owners = Arc::clone(&inner.owners);
511 let logging = LoggingResource::new(Box::new(inner), Arc::new(PanicSink));
512 let budget = super::super::budget::GlobalDeviceBudget::new(Box::new(logging), 64);
513 assert!(std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
514 let _ = budget.allocate(64, StreamId::DEFAULT, AllocTag::UNTAGGED);
515 }))
516 .is_err());
517 crate::cuda_graph::reap_capture_retirements();
518 assert_eq!(ledger.snapshot(), (0, 64));
519 assert_eq!(budget.reserved_bytes(), 0);
520 assert!(owners.lock().unwrap().is_empty());
521 }
522
523 use crate::CudaDevice;
524
525 fn try_device() -> Option<Arc<CudaDevice>> {
526 CudaDevice::new(0).ok().map(Arc::new)
527 }
528
529 #[test]
530 fn pass_through_alloc_dealloc_emits_two_ok_records() {
531 let Some(device) = try_device() else {
532 return;
533 };
534 let inner = Box::new(DirectCudaResource::new(Arc::clone(&device), 0));
535 let sink = Arc::new(InMemorySink::new());
536 let r = LoggingResource::new(inner, sink.clone());
537
538 let block = r
539 .allocate(1024, StreamId::DEFAULT, AllocTag("logging-test"))
540 .expect("alloc");
541 assert_eq!(block.bytes, 1024);
542 assert_eq!(block.state, BlockState::Live);
543 r.deallocate(block).expect("dealloc");
544
545 let recs = sink.snapshot();
546 assert_eq!(recs.len(), 2, "expected 2 records, got {:?}", recs);
547 assert_eq!(recs[0].action, LogAction::Allocate);
548 assert_eq!(recs[0].result, LogResult::Ok);
549 assert_eq!(recs[0].bytes, Some(1024));
550 assert_eq!(recs[0].stream_id, Some(StreamId::DEFAULT));
551 assert!(recs[0].ptr.is_some());
552
553 assert_eq!(recs[1].action, LogAction::Deallocate);
554 assert_eq!(recs[1].result, LogResult::Ok);
555 assert_eq!(recs[1].ptr, recs[0].ptr);
556 assert_eq!(recs[1].generation, recs[0].generation);
557 }
558
559 #[test]
560 fn order_counter_strictly_increases_across_records() {
561 let Some(device) = try_device() else {
562 return;
563 };
564 let inner = Box::new(DirectCudaResource::new(Arc::clone(&device), 0));
565 let sink = Arc::new(InMemorySink::new());
566 let r = LoggingResource::new(inner, sink.clone());
567
568 for _ in 0..4 {
569 let b = r
570 .allocate(64, StreamId::DEFAULT, AllocTag::UNTAGGED)
571 .expect("alloc");
572 r.deallocate(b).expect("dealloc");
573 }
574 r.reap_pending().expect("reap");
575
576 let recs = sink.snapshot();
577 assert_eq!(recs.len(), 9); let mut last = 0u64;
579 for rec in &recs {
580 assert!(
581 rec.order_counter > last,
582 "order_counter must strictly increase: prev={}, now={}",
583 last,
584 rec.order_counter
585 );
586 last = rec.order_counter;
587 }
588 }
589
590 #[test]
591 fn failed_alloc_records_error_result() {
592 let Some(device) = try_device() else {
593 return;
594 };
595 let inner = Box::new(DirectCudaResource::new(Arc::clone(&device), 0));
596 let sink = Arc::new(InMemorySink::new());
597 let r = LoggingResource::new(inner, sink.clone());
598
599 let _ = r.allocate(0, StreamId::DEFAULT, AllocTag::UNTAGGED);
601 let recs = sink.snapshot();
602 assert_eq!(recs.len(), 1);
603 assert_eq!(recs[0].action, LogAction::Allocate);
604 assert!(matches!(recs[0].result, LogResult::Err { kind, .. } if kind == "Driver"));
605 assert_eq!(recs[0].bytes, Some(0));
607 assert!(recs[0].ptr.is_none());
608 assert!(recs[0].generation.is_none());
609 }
610
611 #[test]
612 fn failed_dealloc_records_error_result() {
613 let Some(device) = try_device() else {
614 return;
615 };
616 let inner = Box::new(DirectCudaResource::new(Arc::clone(&device), 0));
617 let sink = Arc::new(InMemorySink::new());
618 let r = LoggingResource::new(inner, sink.clone());
619
620 let bogus = DeviceBlock {
621 ptr: 0xdead_beef,
622 device_ordinal: 0,
623 alloc_stream: StreamId::DEFAULT,
624 bytes: 16,
625 align: 1,
626 tag: AllocTag::UNTAGGED,
627 generation: Generation::next(),
628 state: BlockState::Live,
629 };
630 let res = r.deallocate(bogus);
631 assert!(res.is_err());
632 let recs = sink.snapshot();
633 assert_eq!(recs.len(), 1);
634 assert_eq!(recs[0].action, LogAction::Deallocate);
635 assert!(matches!(recs[0].result, LogResult::Err { kind, .. } if kind == "UseAfterFree"));
636 assert_eq!(recs[0].ptr, Some(0xdead_beef));
637 }
638
639 #[test]
640 fn reap_pending_emits_record() {
641 let Some(device) = try_device() else {
642 return;
643 };
644 let inner = Box::new(DirectCudaResource::new(Arc::clone(&device), 0));
645 let sink = Arc::new(InMemorySink::new());
646 let r = LoggingResource::new(inner, sink.clone());
647
648 r.reap_pending().expect("reap");
649 let recs = sink.snapshot();
650 assert_eq!(recs.len(), 1);
651 assert_eq!(recs[0].action, LogAction::ReapPending);
652 assert_eq!(recs[0].result, LogResult::Ok);
653 assert!(recs[0].stream_id.is_none());
654 assert!(recs[0].ptr.is_none());
655 }
656
657 #[test]
658 fn sink_failure_increments_dropped_records_but_does_not_break_alloc() {
659 struct RefuseAllSink;
662 impl LoggingSink for RefuseAllSink {
663 fn emit(&self, _r: LogRecord) -> Result<(), SinkError> {
664 Err(SinkError::Refused("test sink refuses all".into()))
665 }
666 }
667
668 let Some(device) = try_device() else {
669 return;
670 };
671 let inner = Box::new(DirectCudaResource::new(Arc::clone(&device), 0));
672 let sink = Arc::new(RefuseAllSink);
673 let r = LoggingResource::new(inner, sink);
674
675 let block = r
677 .allocate(128, StreamId::DEFAULT, AllocTag("refuse-test"))
678 .expect("alloc must succeed even when sink refuses");
679 r.deallocate(block).expect("dealloc must succeed too");
680 assert_eq!(r.dropped_records(), 2);
682 }
683
684 #[test]
685 fn forwards_bytes_outstanding_and_device_ordinal() {
686 let Some(device) = try_device() else {
687 return;
688 };
689 let inner = Box::new(DirectCudaResource::new(Arc::clone(&device), 3));
690 let sink = Arc::new(InMemorySink::new());
691 let r = LoggingResource::new(inner, sink);
692
693 assert_eq!(r.device_ordinal(), 3);
694 assert_eq!(r.bytes_outstanding(), 0);
695 }
696
697 #[test]
698 fn truncate_handles_non_ascii_at_cap_without_panicking() {
699 let s = String::from("héllo");
708 let out = truncate(s, 2);
709 assert!(out.starts_with('h'));
710 assert!(out.ends_with('…'));
711
712 let out = truncate(String::from("ok"), 100);
714 assert_eq!(out, "ok");
715
716 let out = truncate(String::from("héllo"), 0);
718 assert_eq!(out, "…");
719
720 let out = truncate(String::from("abcdefgh"), 3);
723 assert_eq!(out, "abc…");
724 }
725
726 #[test]
727 fn null_sink_accepts_records_without_retention() {
728 let sink = NullSink::new();
729 let rec = LogRecord {
730 action: LogAction::Allocate,
731 device_ordinal: 0,
732 stream_id: Some(StreamId::DEFAULT),
733 ptr: Some(0xdead_beef),
734 bytes: Some(64),
735 tag: Some(AllocTag::UNTAGGED),
736 generation: Some(Generation::next()),
737 thread_id: 0,
738 order_counter: 1,
739 timestamp_nanos: 0,
740 result: LogResult::Ok,
741 };
742 sink.emit(rec).expect("NullSink never errors");
743 }
746}