@@ -5,7 +5,7 @@ use core::{
55 convert:: TryInto ,
66 marker:: PhantomData ,
77 mem:: swap,
8- ptr,
8+ ptr:: { self , NonNull } ,
99 sync:: atomic:: { fence, AtomicPtr , AtomicU32 , AtomicU64 , Ordering } ,
1010} ;
1111use std:: {
@@ -68,7 +68,7 @@ static PROCESS_CONTEXT_HANDLER: Mutex<Option<ProcessContextHandle>> = Mutex::new
6868
6969pub ( super ) trait HeaderMemoryHolder : Sized {
7070 fn new ( ) -> io:: Result < Self > ;
71- fn as_ptr ( & self ) -> * mut MappingHeader ;
71+ fn as_ptr ( & self ) -> Option < NonNull < MappingHeader > > ;
7272 fn make_discoverable ( & mut self ) ;
7373 fn unpublish_and_release ( self ) -> io:: Result < ( ) > ;
7474 fn after_fork ( self ) ;
@@ -84,9 +84,6 @@ struct ProcessContextHandleGen<M: HeaderMemoryHolder, T: MonotonicTime> {
8484 /// Once published, and until the next update is complete, the backing allocation of
8585 /// `payload` might be read and thus must not move (e.g. by resizing or drop).
8686 payload : Vec < u8 > ,
87- /// The process id of the last publisher. This is useful to detect forks(), and publish a
88- /// new context accordingly.
89- pid : u32 ,
9087 monotonic_clock : PhantomData < T > ,
9188}
9289
@@ -101,7 +98,11 @@ impl<M: HeaderMemoryHolder, T: MonotonicTime> ProcessContextHandleGen<M, T> {
10198 let mut mapping = M :: new ( ) ?;
10299 let published_at_ns = T :: monotonic_time_ns ( ) ?;
103100
104- let header = mapping. as_ptr ( ) ;
101+ let header = mapping
102+ . as_ptr ( )
103+ // should never happen; as_ptr should only return None after a fork
104+ . ok_or_else ( || io:: Error :: other ( "new process context header mapping is unavailable" ) ) ?
105+ . as_ptr ( ) ;
105106
106107 // SAFETY: header points to a zero-filled, writable allocation of at least
107108 // mapping_size() bytes with MappingHeader alignment; field projections are in-bounds.
@@ -111,12 +112,12 @@ impl<M: HeaderMemoryHolder, T: MonotonicTime> ProcessContextHandleGen<M, T> {
111112 unsafe {
112113 ptr:: addr_of_mut!( ( * header) . signature) . write ( * SIGNATURE ) ;
113114 ptr:: addr_of_mut!( ( * header) . version) . write ( PROCESS_CTX_VERSION ) ;
114- ( * header)
115- . payload_ptr
116- . store ( payload. as_ptr ( ) . cast_mut ( ) , Ordering :: Relaxed ) ;
117115 ( * header)
118116 . payload_size
119117 . store ( payload_size, Ordering :: Relaxed ) ;
118+ ( * header)
119+ . payload_ptr
120+ . store ( payload. as_ptr ( ) . cast_mut ( ) , Ordering :: Relaxed ) ;
120121
121122 fence ( Ordering :: SeqCst ) ;
122123 ( * header)
@@ -129,14 +130,22 @@ impl<M: HeaderMemoryHolder, T: MonotonicTime> ProcessContextHandleGen<M, T> {
129130 Ok ( ProcessContextHandleGen {
130131 mapping,
131132 payload,
132- pid : std:: process:: id ( ) ,
133133 monotonic_clock : PhantomData ,
134134 } )
135135 }
136136
137137 /// Updates the context after initial publication.
138138 fn update ( & mut self , payload : Vec < u8 > ) -> io:: Result < ( ) > {
139- let header = self . mapping . as_ptr ( ) ;
139+ let header = self
140+ . mapping
141+ . as_ptr ( )
142+ . ok_or_else ( || {
143+ io:: Error :: new (
144+ io:: ErrorKind :: InvalidInput ,
145+ "process context header mapping is unavailable after fork" ,
146+ )
147+ } ) ?
148+ . as_ptr ( ) ;
140149
141150 let monotonic_published_at_ns = T :: monotonic_time_ns ( ) ?;
142151 let payload_size: u32 = payload. len ( ) . try_into ( ) . map_err ( |_| {
@@ -157,7 +166,10 @@ impl<M: HeaderMemoryHolder, T: MonotonicTime> ProcessContextHandleGen<M, T> {
157166 . swap ( UNPUBLISHED_OR_UPDATING , Ordering :: Relaxed )
158167 } ;
159168 if last_monotonic_published_at_ns == UNPUBLISHED_OR_UPDATING {
160- panic ! ( "concurrent update of the process context is not supported" ) ;
169+ return Err ( io:: Error :: new (
170+ io:: ErrorKind :: WouldBlock ,
171+ "context is already being updated" ,
172+ ) ) ;
161173 }
162174
163175 let monotonic_published_at_ns =
@@ -172,12 +184,12 @@ impl<M: HeaderMemoryHolder, T: MonotonicTime> ProcessContextHandleGen<M, T> {
172184 self . payload = payload;
173185
174186 unsafe {
175- ( * header)
176- . payload_ptr
177- . store ( self . payload . as_ptr ( ) . cast_mut ( ) , Ordering :: Relaxed ) ;
178187 ( * header)
179188 . payload_size
180189 . store ( payload_size, Ordering :: Relaxed ) ;
190+ ( * header)
191+ . payload_ptr
192+ . store ( self . payload . as_ptr ( ) . cast_mut ( ) , Ordering :: Relaxed ) ;
181193 }
182194
183195 // Prevent the final timestamp publication from moving above either the payload metadata
@@ -214,8 +226,7 @@ fn lock_context_handle() -> io::Result<MutexGuard<'static, Option<ProcessContext
214226///
215227/// - this is the first publication
216228/// - [unpublish] has been called last
217- /// - the previous context has been published from a different process id (that is, a `fork()`
218- /// happened and we're the child process)
229+ /// - the previous header mapping is unavailable after `fork()`
219230///
220231/// Then we follow the Publish protocol of the OTel process context specification (allocating a
221232/// fresh header).
@@ -239,7 +250,7 @@ pub(super) fn publish_raw_payload(payload: Vec<u8>) -> io::Result<()> {
239250 let mut guard = lock_context_handle ( ) ?;
240251
241252 match & mut * guard {
242- Some ( handler) if handler. pid == std :: process :: id ( ) => handler. update ( payload) ,
253+ Some ( handler) if handler. mapping . as_ptr ( ) . is_some ( ) => handler. update ( payload) ,
243254 Some ( handler) => {
244255 let mut local_handler = ProcessContextHandleGen :: publish ( payload) ?;
245256 // If we've been forked, we need to prevent the mapping from being dropped
@@ -272,20 +283,22 @@ pub fn unpublish() -> io::Result<()> {
272283 mapping, payload, ..
273284 } ) = guard. take ( )
274285 {
275- // Mark the context as unavailable before freeing the mapping/payload. The fence forces
276- // the writing CPU not to reorder the unavailable timestamp store and the deallocation
277- // stores. This gives readers more of a chance (but no guarantee) to observe an
278- // unavailable context before the publication is removed.
279- //
280- // SAFETY: the mapping is still live and valid, and the global mutex prevents
281- // concurrent in-process writers from mutating the plain header fields.
282- let header = mapping. as_ptr ( ) ;
283- unsafe {
284- ( * header)
285- . monotonic_published_at_ns
286- . store ( UNPUBLISHED_OR_UPDATING , Ordering :: Relaxed ) ;
286+ if let Some ( header) = mapping. as_ptr ( ) {
287+ // Mark the context as unavailable before freeing the mapping/payload. The fence
288+ // forces the writing CPU not to reorder the unavailable timestamp store and the
289+ // deallocation stores. This gives readers more of a chance (but no guarantee) to
290+ // observe an unavailable context before the publication is removed.
291+ //
292+ // SAFETY: the mapping is still live and valid, and the global mutex prevents
293+ // concurrent in-process writers from mutating the plain header fields.
294+ let header = header. as_ptr ( ) ;
295+ unsafe {
296+ ( * header)
297+ . monotonic_published_at_ns
298+ . store ( UNPUBLISHED_OR_UPDATING , Ordering :: Relaxed ) ;
299+ }
300+ fence ( Ordering :: SeqCst ) ;
287301 }
288- fence ( Ordering :: SeqCst ) ;
289302
290303 // The payload will still drop if this fails, leaving a zero timestamp behind.
291304 mapping. unpublish_and_release ( ) ?;
0 commit comments