Skip to main content

heapdump_snapshot/
snapshot.rs

1// Copyright 2023 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use fidl::MonotonicInstant;
6use flex_fuchsia_memory_heapdump_client as fheapdump_client;
7use futures::StreamExt;
8use std::collections::{HashMap, HashSet};
9use std::rc::Rc;
10
11use crate::Error;
12
13/// Contains a snapshot along with metadata from its header.
14#[derive(Debug)]
15pub struct SnapshotWithHeader {
16    pub process_name: String,
17    pub process_koid: u64,
18    pub snapshot: Snapshot,
19}
20
21/// Contains all the data received over a `SnapshotReceiver` channel.
22#[derive(Debug)]
23pub struct Snapshot {
24    /// All the live allocations in the analyzed process, indexed by memory address.
25    pub allocations: Vec<Allocation>,
26
27    /// All the executable memory regions in the analyzed process, indexed by start address.
28    pub executable_regions: HashMap<u64, ExecutableRegion>,
29}
30
31/// Information about one or more allocated memory blocks.
32#[derive(Debug)]
33pub struct Allocation {
34    pub address: Option<u64>,
35
36    /// Number of allocations that have been aggregated into this `Allocation` instance.
37    pub count: u64,
38
39    /// Total size, in bytes.
40    pub size: u64,
41
42    /// Allocating thread.
43    pub thread_info: Option<Rc<ThreadInfo>>,
44
45    /// Stack trace of the allocation site.
46    pub stack_trace: Rc<StackTrace>,
47
48    /// Allocation timestamp, in nanoseconds.
49    pub timestamp: Option<MonotonicInstant>,
50
51    /// Memory dump of this block's contents.
52    pub contents: Option<Vec<u8>>,
53}
54
55/// A stack trace.
56#[derive(Debug)]
57pub struct StackTrace {
58    /// Code addresses at each call frame. The first entry corresponds to the leaf call.
59    pub program_addresses: Vec<u64>,
60}
61
62/// A memory region containing code loaded from an ELF file.
63#[derive(Debug)]
64pub struct ExecutableRegion {
65    /// Region name for human consumption (usually either the ELF soname or the VMO name), if known.
66    pub name: String,
67
68    /// Region size, in bytes.
69    pub size: u64,
70
71    /// The corresponding offset in the ELF file.
72    pub file_offset: u64,
73
74    /// The corresponding relative address in the ELF file.
75    pub vaddr: u64,
76
77    /// The Build ID of the ELF file.
78    pub build_id: Vec<u8>,
79}
80
81/// Information identifying a specific thread.
82#[derive(Debug, PartialEq)]
83pub struct ThreadInfo {
84    /// The thread's koid.
85    pub koid: zx_types::zx_koid_t,
86
87    /// The thread's name.
88    pub name: String,
89}
90
91/// Gets the value of a field in a FIDL table as a `Result<T, Error>`.
92///
93/// An `Err(Error::MissingField { .. })` is returned if the field's value is `None`.
94///
95/// Usage: `read_field!(container_expression => ContainerType, field_name)`
96///
97/// # Example
98///
99/// ```
100/// struct MyFidlTable { field: Option<u32>, .. }
101/// let table = MyFidlTable { field: Some(44), .. };
102///
103/// let val = read_field!(table => MyFidlTable, field)?;
104/// ```
105macro_rules! read_field {
106    ($e:expr => $c:ident, $f:ident) => {
107        $e.$f.ok_or(Error::MissingField {
108            container: std::stringify!($c),
109            field: std::stringify!($f),
110        })
111    };
112}
113
114impl Snapshot {
115    /// Receives a snapshot over a `SnapshotReceiver` channel and reassembles it.
116    pub async fn receive_single_from(
117        mut stream: fheapdump_client::SnapshotReceiverRequestStream,
118    ) -> Result<Snapshot, Error> {
119        Snapshot::receive_inner(&mut stream).await
120    }
121
122    /// Receives multiple header-prefixed snapshots over a `SnapshotReceiver` channel and
123    /// reassemble them.
124    #[cfg(fuchsia_api_level_at_least = "HEAD")]
125    pub async fn receive_multi_from(
126        mut stream: fheapdump_client::SnapshotReceiverRequestStream,
127    ) -> Result<Vec<SnapshotWithHeader>, Error> {
128        let mut snapshots = vec![];
129        loop {
130            // Wait for a batch of elements containing either just a header element or an empty
131            // batch (to signal the end of the stream).
132            match stream.next().await.transpose()? {
133                Some(fheapdump_client::SnapshotReceiverRequest::Batch { batch, responder }) => {
134                    match &batch[..] {
135                        [fheapdump_client::SnapshotElement::SnapshotHeader(header)] => {
136                            responder.send()?;
137
138                            // Receive the actual snapshot.
139                            let snapshot = Snapshot::receive_inner(&mut stream).await?;
140
141                            let header = header.clone();
142                            snapshots.push(SnapshotWithHeader {
143                                process_name: read_field!(header => SnapshotHeader, process_name)?,
144                                process_koid: read_field!(header => SnapshotHeader, process_koid)?,
145                                snapshot,
146                            });
147                        }
148                        [] => {
149                            responder.send()?;
150                            return Ok(snapshots);
151                        }
152                        _ => return Err(Error::HeaderExpected),
153                    }
154                }
155                Some(fheapdump_client::SnapshotReceiverRequest::ReportError {
156                    error,
157                    responder,
158                }) => {
159                    let _ = responder.send(); // Ignore the result of the acknowledgment.
160                    return Err(Error::CollectorError(error));
161                }
162                None => return Err(Error::UnexpectedEndOfStream),
163            };
164        }
165    }
166
167    async fn receive_inner(
168        stream: &mut fheapdump_client::SnapshotReceiverRequestStream,
169    ) -> Result<Snapshot, Error> {
170        struct AllocationValue {
171            address: Option<u64>,
172            count: u64,
173            size: u64,
174            thread_info_key: Option<u64>,
175            stack_trace_key: u64,
176            timestamp: Option<MonotonicInstant>,
177        }
178        let mut allocation_addresses: HashSet<u64> = HashSet::new();
179        let mut allocations: Vec<AllocationValue> = vec![];
180        let mut thread_infos: HashMap<u64, Rc<ThreadInfo>> = HashMap::new();
181        let mut stack_traces: HashMap<u64, Vec<u64>> = HashMap::new();
182        let mut executable_regions: HashMap<u64, ExecutableRegion> = HashMap::new();
183        let mut contents: HashMap<u64, Vec<u8>> = HashMap::new();
184
185        loop {
186            // Wait for the next batch of elements.
187            let batch = match stream.next().await.transpose()? {
188                Some(fheapdump_client::SnapshotReceiverRequest::Batch { batch, responder }) => {
189                    // Send acknowledgment as quickly as possible, then keep processing the received batch.
190                    responder.send()?;
191                    batch
192                }
193                Some(fheapdump_client::SnapshotReceiverRequest::ReportError {
194                    error,
195                    responder,
196                }) => {
197                    let _ = responder.send(); // Ignore the result of the acknowledgment.
198                    return Err(Error::CollectorError(error));
199                }
200                None => return Err(Error::UnexpectedEndOfStream),
201            };
202
203            // Process data. An empty batch signals the end of the stream.
204            if !batch.is_empty() {
205                for element in batch {
206                    match element {
207                        fheapdump_client::SnapshotElement::Allocation(allocation) => {
208                            if let Some(address) = allocation.address {
209                                if !allocation_addresses.insert(address) {
210                                    return Err(Error::ConflictingElement {
211                                        element_type: "Allocation",
212                                    });
213                                }
214                            }
215
216                            #[cfg(not(fuchsia_api_level_at_least = "29"))]
217                            let count = 1;
218                            #[cfg(fuchsia_api_level_at_least = "29")]
219                            let count = allocation.count.unwrap_or(1);
220
221                            let size = read_field!(allocation => Allocation, size)?;
222                            let stack_trace_key =
223                                read_field!(allocation => Allocation, stack_trace_key)?;
224                            allocations.push(AllocationValue {
225                                address: allocation.address,
226                                count,
227                                size,
228                                thread_info_key: allocation.thread_info_key,
229                                stack_trace_key,
230                                timestamp: allocation.timestamp,
231                            });
232                        }
233                        fheapdump_client::SnapshotElement::StackTrace(stack_trace) => {
234                            let stack_trace_key =
235                                read_field!(stack_trace => StackTrace, stack_trace_key)?;
236                            let mut program_addresses =
237                                read_field!(stack_trace => StackTrace, program_addresses)?;
238                            stack_traces
239                                .entry(stack_trace_key)
240                                .or_default()
241                                .append(&mut program_addresses);
242                        }
243                        fheapdump_client::SnapshotElement::ThreadInfo(thread_info) => {
244                            let thread_info_key =
245                                read_field!(thread_info => ThreadInfo, thread_info_key)?;
246                            let koid = read_field!(thread_info => ThreadInfo, koid)?;
247                            let name = read_field!(thread_info => ThreadInfo, name)?;
248                            if thread_infos
249                                .insert(thread_info_key, Rc::new(ThreadInfo { koid, name }))
250                                .is_some()
251                            {
252                                return Err(Error::ConflictingElement {
253                                    element_type: "ThreadInfo",
254                                });
255                            }
256                        }
257                        fheapdump_client::SnapshotElement::ExecutableRegion(region) => {
258                            let address = read_field!(region => ExecutableRegion, address)?;
259                            let name = region.name.unwrap_or_else(|| String::new());
260                            let size = read_field!(region => ExecutableRegion, size)?;
261                            let file_offset = read_field!(region => ExecutableRegion, file_offset)?;
262                            let vaddr = read_field!(region => ExecutableRegion, vaddr)?;
263                            let build_id = read_field!(region => ExecutableRegion, build_id)?.value;
264                            let region =
265                                ExecutableRegion { name, size, file_offset, vaddr, build_id };
266                            if executable_regions.insert(address, region).is_some() {
267                                return Err(Error::ConflictingElement {
268                                    element_type: "ExecutableRegion",
269                                });
270                            }
271                        }
272                        fheapdump_client::SnapshotElement::BlockContents(block_contents) => {
273                            let address = read_field!(block_contents => BlockContents, address)?;
274                            let mut chunk = read_field!(block_contents => BlockContents, contents)?;
275                            contents.entry(address).or_default().append(&mut chunk);
276                        }
277                        _ => return Err(Error::UnexpectedElementType),
278                    }
279                }
280            } else {
281                // We are at the end of the stream. Convert to the final types and resolve
282                // cross-references.
283                let final_stack_traces: HashMap<u64, Rc<StackTrace>> = stack_traces
284                    .into_iter()
285                    .map(|(key, program_addresses)| {
286                        (key, Rc::new(StackTrace { program_addresses }))
287                    })
288                    .collect();
289                let mut final_allocations = vec![];
290                for AllocationValue {
291                    address,
292                    count,
293                    size,
294                    thread_info_key,
295                    stack_trace_key,
296                    timestamp,
297                } in allocations
298                {
299                    let thread_info = match thread_info_key {
300                        Some(key) => Some(
301                            thread_infos
302                                .get(&key)
303                                .ok_or(Error::InvalidCrossReference { element_type: "ThreadInfo" })?
304                                .clone(),
305                        ),
306                        None => None,
307                    };
308                    let stack_trace = final_stack_traces
309                        .get(&stack_trace_key)
310                        .ok_or(Error::InvalidCrossReference { element_type: "StackTrace" })?
311                        .clone();
312                    let contents = address.and_then(|address| contents.remove(&address));
313                    if let Some(data) = &contents {
314                        if data.len() as u64 != size {
315                            return Err(Error::ConflictingElement {
316                                element_type: "BlockContents",
317                            });
318                        }
319                    }
320                    final_allocations.push(Allocation {
321                        address,
322                        count,
323                        size,
324                        thread_info,
325                        stack_trace,
326                        timestamp,
327                        contents,
328                    });
329                }
330
331                return Ok(Snapshot { allocations: final_allocations, executable_regions });
332            }
333        }
334    }
335}
336
337#[cfg(test)]
338mod tests {
339    use super::*;
340    use crate::test_helpers::create_client;
341    use assert_matches::assert_matches;
342    use fuchsia_async as fasync;
343    use test_case::test_case;
344
345    // Constants used by some of the tests below:
346    const FAKE_ALLOCATION_1_ADDRESS: u64 = 1234;
347    const FAKE_ALLOCATION_1_SIZE: u64 = 8;
348    const FAKE_ALLOCATION_1_TIMESTAMP: MonotonicInstant = MonotonicInstant::from_nanos(888888888);
349    const FAKE_ALLOCATION_1_CONTENTS: [u8; FAKE_ALLOCATION_1_SIZE as usize] = *b"12345678";
350    const FAKE_ALLOCATION_2_ADDRESS: u64 = 5678;
351    const FAKE_ALLOCATION_2_SIZE: u64 = 4;
352    const FAKE_ALLOCATION_2_TIMESTAMP: MonotonicInstant = MonotonicInstant::from_nanos(-777777777); // test negative value too
353    const FAKE_THREAD_1_KOID: u64 = 1212;
354    const FAKE_THREAD_1_NAME: &str = "fake-thread-1-name";
355    const FAKE_THREAD_1_KEY: u64 = 4567;
356    const FAKE_THREAD_2_KOID: u64 = 1213;
357    const FAKE_THREAD_2_NAME: &str = "fake-thread-2-name";
358    const FAKE_THREAD_2_KEY: u64 = 7654;
359    const FAKE_STACK_TRACE_1_ADDRESSES: [u64; 6] = [11111, 22222, 33333, 22222, 44444, 55555];
360    const FAKE_STACK_TRACE_1_KEY: u64 = 9876;
361    const FAKE_STACK_TRACE_2_ADDRESSES: [u64; 4] = [11111, 22222, 11111, 66666];
362    const FAKE_STACK_TRACE_2_KEY: u64 = 6789;
363    const FAKE_REGION_1_ADDRESS: u64 = 0x10000000;
364    const FAKE_REGION_1_NAME: &str = "region-1";
365    const FAKE_REGION_1_SIZE: u64 = 0x80000;
366    const FAKE_REGION_1_FILE_OFFSET: u64 = 0x1000;
367    const FAKE_REGION_1_VADDR: u64 = 0x3000;
368    const FAKE_REGION_1_BUILD_ID: &[u8] = &[0xaa; 20];
369    const FAKE_REGION_2_ADDRESS: u64 = 0x7654300000;
370    const FAKE_REGION_2_SIZE: u64 = 0x200000;
371    const FAKE_REGION_2_FILE_OFFSET: u64 = 0x2000;
372    const FAKE_REGION_2_BUILD_ID: &[u8] = &[0x55; 32];
373    const FAKE_REGION_2_VADDR: u64 = 0x7000;
374
375    #[fasync::run_singlethreaded(test)]
376    async fn test_empty() {
377        let client = create_client();
378        let (receiver_proxy, receiver_stream) =
379            client.create_proxy_and_stream::<fheapdump_client::SnapshotReceiverMarker>();
380        let receive_worker = fasync::Task::local(Snapshot::receive_single_from(receiver_stream));
381
382        // Send the end of stream marker.
383        let fut = receiver_proxy.batch(&[]);
384        fut.await.unwrap();
385
386        // Receive the snapshot we just transmitted and verify that it is empty.
387        let received_snapshot = receive_worker.await.unwrap();
388        assert!(received_snapshot.allocations.is_empty());
389        assert!(received_snapshot.executable_regions.is_empty());
390    }
391
392    #[fasync::run_singlethreaded(test)]
393    async fn test_one_batch() {
394        let client = create_client();
395        let (receiver_proxy, receiver_stream) =
396            client.create_proxy_and_stream::<fheapdump_client::SnapshotReceiverMarker>();
397        let receive_worker = fasync::Task::local(Snapshot::receive_single_from(receiver_stream));
398
399        // Send a batch containing two allocations - whose threads, stack traces and contents can be
400        // listed before or after the allocation(s) that reference them - and two executable
401        // regions.
402        let fut = receiver_proxy.batch(&[
403            fheapdump_client::SnapshotElement::BlockContents(fheapdump_client::BlockContents {
404                address: Some(FAKE_ALLOCATION_1_ADDRESS),
405                contents: Some(FAKE_ALLOCATION_1_CONTENTS.to_vec()),
406                ..Default::default()
407            }),
408            fheapdump_client::SnapshotElement::ExecutableRegion(
409                fheapdump_client::ExecutableRegion {
410                    address: Some(FAKE_REGION_1_ADDRESS),
411                    name: Some(FAKE_REGION_1_NAME.to_string()),
412                    size: Some(FAKE_REGION_1_SIZE),
413                    file_offset: Some(FAKE_REGION_1_FILE_OFFSET),
414                    vaddr: Some(FAKE_REGION_1_VADDR),
415                    build_id: Some(fheapdump_client::BuildId {
416                        value: FAKE_REGION_1_BUILD_ID.to_vec(),
417                    }),
418                    ..Default::default()
419                },
420            ),
421            fheapdump_client::SnapshotElement::StackTrace(fheapdump_client::StackTrace {
422                stack_trace_key: Some(FAKE_STACK_TRACE_1_KEY),
423                program_addresses: Some(FAKE_STACK_TRACE_1_ADDRESSES.to_vec()),
424                ..Default::default()
425            }),
426            fheapdump_client::SnapshotElement::Allocation(fheapdump_client::Allocation {
427                address: Some(FAKE_ALLOCATION_1_ADDRESS),
428                size: Some(FAKE_ALLOCATION_1_SIZE),
429                thread_info_key: Some(FAKE_THREAD_1_KEY),
430                stack_trace_key: Some(FAKE_STACK_TRACE_2_KEY),
431                timestamp: Some(FAKE_ALLOCATION_1_TIMESTAMP),
432                ..Default::default()
433            }),
434            fheapdump_client::SnapshotElement::ThreadInfo(fheapdump_client::ThreadInfo {
435                thread_info_key: Some(FAKE_THREAD_1_KEY),
436                koid: Some(FAKE_THREAD_1_KOID),
437                name: Some(FAKE_THREAD_1_NAME.to_string()),
438                ..Default::default()
439            }),
440            fheapdump_client::SnapshotElement::ExecutableRegion(
441                fheapdump_client::ExecutableRegion {
442                    address: Some(FAKE_REGION_2_ADDRESS),
443                    size: Some(FAKE_REGION_2_SIZE),
444                    file_offset: Some(FAKE_REGION_2_FILE_OFFSET),
445                    build_id: Some(fheapdump_client::BuildId {
446                        value: FAKE_REGION_2_BUILD_ID.to_vec(),
447                    }),
448                    vaddr: Some(FAKE_REGION_2_VADDR),
449                    ..Default::default()
450                },
451            ),
452            fheapdump_client::SnapshotElement::ThreadInfo(fheapdump_client::ThreadInfo {
453                thread_info_key: Some(FAKE_THREAD_2_KEY),
454                koid: Some(FAKE_THREAD_2_KOID),
455                name: Some(FAKE_THREAD_2_NAME.to_string()),
456                ..Default::default()
457            }),
458            fheapdump_client::SnapshotElement::Allocation(fheapdump_client::Allocation {
459                address: Some(FAKE_ALLOCATION_2_ADDRESS),
460                size: Some(FAKE_ALLOCATION_2_SIZE),
461                thread_info_key: Some(FAKE_THREAD_2_KEY),
462                stack_trace_key: Some(FAKE_STACK_TRACE_1_KEY),
463                timestamp: Some(FAKE_ALLOCATION_2_TIMESTAMP),
464                ..Default::default()
465            }),
466            fheapdump_client::SnapshotElement::StackTrace(fheapdump_client::StackTrace {
467                stack_trace_key: Some(FAKE_STACK_TRACE_2_KEY),
468                program_addresses: Some(FAKE_STACK_TRACE_2_ADDRESSES.to_vec()),
469                ..Default::default()
470            }),
471        ]);
472        fut.await.unwrap();
473
474        // Send the end of stream marker.
475        let fut = receiver_proxy.batch(&[]);
476        fut.await.unwrap();
477
478        // Receive the snapshot we just transmitted and verify its contents.
479        let mut received_snapshot = receive_worker.await.unwrap();
480        let allocation1 = received_snapshot.allocations.swap_remove(
481            received_snapshot
482                .allocations
483                .iter()
484                .position(|alloc| alloc.address == Some(FAKE_ALLOCATION_1_ADDRESS))
485                .unwrap(),
486        );
487        assert_eq!(allocation1.size, FAKE_ALLOCATION_1_SIZE);
488        assert_eq!(
489            allocation1.thread_info,
490            Some(Rc::new(ThreadInfo {
491                koid: FAKE_THREAD_1_KOID,
492                name: FAKE_THREAD_1_NAME.to_owned()
493            }))
494        );
495        assert_eq!(allocation1.stack_trace.program_addresses, FAKE_STACK_TRACE_2_ADDRESSES);
496        assert_eq!(allocation1.timestamp, Some(FAKE_ALLOCATION_1_TIMESTAMP));
497        assert_eq!(
498            allocation1.contents.as_ref().expect("contents must be set"),
499            &FAKE_ALLOCATION_1_CONTENTS.to_vec()
500        );
501        let allocation2 = received_snapshot.allocations.swap_remove(
502            received_snapshot
503                .allocations
504                .iter()
505                .position(|alloc| alloc.address == Some(FAKE_ALLOCATION_2_ADDRESS))
506                .unwrap(),
507        );
508        assert_eq!(allocation2.size, FAKE_ALLOCATION_2_SIZE);
509        assert_eq!(
510            allocation2.thread_info,
511            Some(Rc::new(ThreadInfo {
512                koid: FAKE_THREAD_2_KOID,
513                name: FAKE_THREAD_2_NAME.to_owned()
514            }))
515        );
516        assert_eq!(allocation2.stack_trace.program_addresses, FAKE_STACK_TRACE_1_ADDRESSES);
517        assert_eq!(allocation2.timestamp, Some(FAKE_ALLOCATION_2_TIMESTAMP));
518        assert_matches!(allocation2.contents, None, "no contents are sent for this allocation");
519        assert!(received_snapshot.allocations.is_empty(), "all the entries have been removed");
520        let region1 = received_snapshot.executable_regions.remove(&FAKE_REGION_1_ADDRESS).unwrap();
521        assert_eq!(region1.name, FAKE_REGION_1_NAME);
522        assert_eq!(region1.size, FAKE_REGION_1_SIZE);
523        assert_eq!(region1.file_offset, FAKE_REGION_1_FILE_OFFSET);
524        assert_eq!(region1.vaddr, FAKE_REGION_1_VADDR);
525        assert_eq!(region1.build_id, FAKE_REGION_1_BUILD_ID);
526        let region2 = received_snapshot.executable_regions.remove(&FAKE_REGION_2_ADDRESS).unwrap();
527        assert_eq!(region2.size, FAKE_REGION_2_SIZE);
528        assert_eq!(region2.file_offset, FAKE_REGION_2_FILE_OFFSET);
529        assert_eq!(region2.build_id, FAKE_REGION_2_BUILD_ID);
530        assert!(received_snapshot.executable_regions.is_empty(), "all entries have been removed");
531    }
532
533    #[fasync::run_singlethreaded(test)]
534    async fn test_two_batches() {
535        let client = create_client();
536        let (receiver_proxy, receiver_stream) =
537            client.create_proxy_and_stream::<fheapdump_client::SnapshotReceiverMarker>();
538        let receive_worker = fasync::Task::local(Snapshot::receive_single_from(receiver_stream));
539
540        // Send a first batch.
541        let fut = receiver_proxy.batch(&[
542            fheapdump_client::SnapshotElement::ExecutableRegion(
543                fheapdump_client::ExecutableRegion {
544                    address: Some(FAKE_REGION_2_ADDRESS),
545                    size: Some(FAKE_REGION_2_SIZE),
546                    file_offset: Some(FAKE_REGION_2_FILE_OFFSET),
547                    build_id: Some(fheapdump_client::BuildId {
548                        value: FAKE_REGION_2_BUILD_ID.to_vec(),
549                    }),
550                    vaddr: Some(FAKE_REGION_2_VADDR),
551                    ..Default::default()
552                },
553            ),
554            fheapdump_client::SnapshotElement::Allocation(fheapdump_client::Allocation {
555                address: Some(FAKE_ALLOCATION_1_ADDRESS),
556                size: Some(FAKE_ALLOCATION_1_SIZE),
557                thread_info_key: Some(FAKE_THREAD_1_KEY),
558                stack_trace_key: Some(FAKE_STACK_TRACE_2_KEY),
559                timestamp: Some(FAKE_ALLOCATION_1_TIMESTAMP),
560                ..Default::default()
561            }),
562            fheapdump_client::SnapshotElement::StackTrace(fheapdump_client::StackTrace {
563                stack_trace_key: Some(FAKE_STACK_TRACE_1_KEY),
564                program_addresses: Some(FAKE_STACK_TRACE_1_ADDRESSES.to_vec()),
565                ..Default::default()
566            }),
567            fheapdump_client::SnapshotElement::ThreadInfo(fheapdump_client::ThreadInfo {
568                thread_info_key: Some(FAKE_THREAD_2_KEY),
569                koid: Some(FAKE_THREAD_2_KOID),
570                name: Some(FAKE_THREAD_2_NAME.to_string()),
571                ..Default::default()
572            }),
573        ]);
574        fut.await.unwrap();
575
576        // Send another batch.
577        let fut = receiver_proxy.batch(&[
578            fheapdump_client::SnapshotElement::ThreadInfo(fheapdump_client::ThreadInfo {
579                thread_info_key: Some(FAKE_THREAD_1_KEY),
580                koid: Some(FAKE_THREAD_1_KOID),
581                name: Some(FAKE_THREAD_1_NAME.to_string()),
582                ..Default::default()
583            }),
584            fheapdump_client::SnapshotElement::Allocation(fheapdump_client::Allocation {
585                address: Some(FAKE_ALLOCATION_2_ADDRESS),
586                size: Some(FAKE_ALLOCATION_2_SIZE),
587                thread_info_key: Some(FAKE_THREAD_2_KEY),
588                stack_trace_key: Some(FAKE_STACK_TRACE_1_KEY),
589                timestamp: Some(FAKE_ALLOCATION_2_TIMESTAMP),
590                ..Default::default()
591            }),
592            fheapdump_client::SnapshotElement::ExecutableRegion(
593                fheapdump_client::ExecutableRegion {
594                    address: Some(FAKE_REGION_1_ADDRESS),
595                    name: Some(FAKE_REGION_1_NAME.to_string()),
596                    size: Some(FAKE_REGION_1_SIZE),
597                    file_offset: Some(FAKE_REGION_1_FILE_OFFSET),
598                    vaddr: Some(FAKE_REGION_1_VADDR),
599                    build_id: Some(fheapdump_client::BuildId {
600                        value: FAKE_REGION_1_BUILD_ID.to_vec(),
601                    }),
602                    ..Default::default()
603                },
604            ),
605            fheapdump_client::SnapshotElement::StackTrace(fheapdump_client::StackTrace {
606                stack_trace_key: Some(FAKE_STACK_TRACE_2_KEY),
607                program_addresses: Some(FAKE_STACK_TRACE_2_ADDRESSES.to_vec()),
608                ..Default::default()
609            }),
610            fheapdump_client::SnapshotElement::BlockContents(fheapdump_client::BlockContents {
611                address: Some(FAKE_ALLOCATION_1_ADDRESS),
612                contents: Some(FAKE_ALLOCATION_1_CONTENTS.to_vec()),
613                ..Default::default()
614            }),
615        ]);
616        fut.await.unwrap();
617
618        // Send the end of stream marker.
619        let fut = receiver_proxy.batch(&[]);
620        fut.await.unwrap();
621
622        // Receive the snapshot we just transmitted and verify its contents.
623        let mut received_snapshot = receive_worker.await.unwrap();
624        let allocation1 = received_snapshot.allocations.swap_remove(
625            received_snapshot
626                .allocations
627                .iter()
628                .position(|alloc| alloc.address == Some(FAKE_ALLOCATION_1_ADDRESS))
629                .unwrap(),
630        );
631        assert_eq!(allocation1.size, FAKE_ALLOCATION_1_SIZE);
632        assert_eq!(
633            allocation1.thread_info,
634            Some(Rc::new(ThreadInfo {
635                koid: FAKE_THREAD_1_KOID,
636                name: FAKE_THREAD_1_NAME.to_owned()
637            }))
638        );
639        assert_eq!(allocation1.stack_trace.program_addresses, FAKE_STACK_TRACE_2_ADDRESSES);
640        assert_eq!(allocation1.timestamp, Some(FAKE_ALLOCATION_1_TIMESTAMP));
641        assert_eq!(
642            allocation1.contents.as_ref().expect("contents must be set"),
643            &FAKE_ALLOCATION_1_CONTENTS.to_vec()
644        );
645        let allocation2 = received_snapshot.allocations.swap_remove(
646            received_snapshot
647                .allocations
648                .iter()
649                .position(|alloc| alloc.address == Some(FAKE_ALLOCATION_2_ADDRESS))
650                .unwrap(),
651        );
652        assert_eq!(allocation2.size, FAKE_ALLOCATION_2_SIZE);
653        assert_eq!(
654            allocation2.thread_info,
655            Some(Rc::new(ThreadInfo {
656                koid: FAKE_THREAD_2_KOID,
657                name: FAKE_THREAD_2_NAME.to_owned()
658            }))
659        );
660        assert_eq!(allocation2.stack_trace.program_addresses, FAKE_STACK_TRACE_1_ADDRESSES);
661        assert_eq!(allocation2.timestamp, Some(FAKE_ALLOCATION_2_TIMESTAMP));
662        assert_matches!(allocation2.contents, None, "no contents are sent for this allocation");
663        assert!(received_snapshot.allocations.is_empty(), "all the entries have been removed");
664        let region1 = received_snapshot.executable_regions.remove(&FAKE_REGION_1_ADDRESS).unwrap();
665        assert_eq!(region1.name, FAKE_REGION_1_NAME);
666        assert_eq!(region1.size, FAKE_REGION_1_SIZE);
667        assert_eq!(region1.file_offset, FAKE_REGION_1_FILE_OFFSET);
668        assert_eq!(region1.vaddr, FAKE_REGION_1_VADDR);
669        assert_eq!(region1.build_id, FAKE_REGION_1_BUILD_ID);
670        let region2 = received_snapshot.executable_regions.remove(&FAKE_REGION_2_ADDRESS).unwrap();
671        assert_eq!(region2.size, FAKE_REGION_2_SIZE);
672        assert_eq!(region2.file_offset, FAKE_REGION_2_FILE_OFFSET);
673        assert_eq!(region2.build_id, FAKE_REGION_2_BUILD_ID);
674        assert!(received_snapshot.executable_regions.is_empty(), "all entries have been removed");
675    }
676
677    #[test_case(|allocation| allocation.size = None => matches
678        Err(Error::MissingField { container: "Allocation", field: "size" }) ; "size")]
679    #[test_case(|allocation| allocation.stack_trace_key = None => matches
680        Err(Error::MissingField { container: "Allocation", field: "stack_trace_key" }) ; "stack_trace_key")]
681    #[test_case(|allocation| allocation.address = None => matches
682        Ok(_) ; "address_is_optional")]
683    #[test_case(|allocation| allocation.thread_info_key = None => matches
684        Ok(_) ; "thread_info_is_optional")]
685    #[test_case(|allocation| allocation.timestamp = None => matches
686        Ok(_) ; "timestamp_is_optional")]
687    #[test_case(|_| () /* if we do not set any field to None, the result should be Ok */ => matches
688        Ok(_) ; "success")]
689    #[fasync::run_singlethreaded(test)]
690    async fn test_allocation_required_fields(
691        set_one_field_to_none: fn(&mut fheapdump_client::Allocation),
692    ) -> Result<Snapshot, Error> {
693        let client = create_client();
694        let (receiver_proxy, receiver_stream) =
695            client.create_proxy_and_stream::<fheapdump_client::SnapshotReceiverMarker>();
696        let receive_worker = fasync::Task::local(Snapshot::receive_single_from(receiver_stream));
697
698        // Start with an Allocation with all the required fields set.
699        let mut allocation = fheapdump_client::Allocation {
700            address: Some(FAKE_ALLOCATION_1_ADDRESS),
701            size: Some(FAKE_ALLOCATION_1_SIZE),
702            thread_info_key: Some(FAKE_THREAD_1_KEY),
703            stack_trace_key: Some(FAKE_STACK_TRACE_1_KEY),
704            timestamp: Some(FAKE_ALLOCATION_1_TIMESTAMP),
705            ..Default::default()
706        };
707
708        // Set one of the fields to None, according to the case being tested.
709        set_one_field_to_none(&mut allocation);
710
711        // Send it to the SnapshotReceiver along with the thread info and stack trace it references.
712        let fut = receiver_proxy.batch(&[
713            fheapdump_client::SnapshotElement::Allocation(allocation),
714            fheapdump_client::SnapshotElement::ThreadInfo(fheapdump_client::ThreadInfo {
715                thread_info_key: Some(FAKE_THREAD_1_KEY),
716                koid: Some(FAKE_THREAD_1_KOID),
717                name: Some(FAKE_THREAD_1_NAME.to_string()),
718                ..Default::default()
719            }),
720            fheapdump_client::SnapshotElement::StackTrace(fheapdump_client::StackTrace {
721                stack_trace_key: Some(FAKE_STACK_TRACE_1_KEY),
722                program_addresses: Some(FAKE_STACK_TRACE_1_ADDRESSES.to_vec()),
723                ..Default::default()
724            }),
725        ]);
726        let _ = fut.await; // ignore result, as the peer may detect the error and close the channel
727
728        // Send the end of stream marker.
729        let fut = receiver_proxy.batch(&[]);
730        let _ = fut.await; // ignore result, as the peer may detect the error and close the channel
731
732        // Return the result.
733        receive_worker.await
734    }
735
736    #[test_case(|thread_info| thread_info.thread_info_key = None => matches
737        Err(Error::MissingField { container: "ThreadInfo", field: "thread_info_key" }) ; "thread_info_key")]
738    #[test_case(|thread_info| thread_info.koid = None => matches
739        Err(Error::MissingField { container: "ThreadInfo", field: "koid" }) ; "koid")]
740    #[test_case(|thread_info| thread_info.name = None => matches
741        Err(Error::MissingField { container: "ThreadInfo", field: "name" }) ; "name")]
742    #[test_case(|_| () /* if we do not set any field to None, the result should be Ok */ => matches
743        Ok(_) ; "success")]
744    #[fasync::run_singlethreaded(test)]
745    async fn test_thread_info_required_fields(
746        set_one_field_to_none: fn(&mut fheapdump_client::ThreadInfo),
747    ) -> Result<Snapshot, Error> {
748        let client = create_client();
749        let (receiver_proxy, receiver_stream) =
750            client.create_proxy_and_stream::<fheapdump_client::SnapshotReceiverMarker>();
751        let receive_worker = fasync::Task::local(Snapshot::receive_single_from(receiver_stream));
752
753        // Start with a ThreadInfo with all the required fields set.
754        let mut thread_info = fheapdump_client::ThreadInfo {
755            thread_info_key: Some(FAKE_THREAD_1_KEY),
756            koid: Some(FAKE_THREAD_1_KOID),
757            name: Some(FAKE_THREAD_1_NAME.to_string()),
758            ..Default::default()
759        };
760
761        // Set one of the fields to None, according to the case being tested.
762        set_one_field_to_none(&mut thread_info);
763
764        // Send it to the SnapshotReceiver.
765        let fut =
766            receiver_proxy.batch(&[fheapdump_client::SnapshotElement::ThreadInfo(thread_info)]);
767        let _ = fut.await; // ignore result, as the peer may detect the error and close the channel
768
769        // Send the end of stream marker.
770        let fut = receiver_proxy.batch(&[]);
771        let _ = fut.await; // ignore result, as the peer may detect the error and close the channel
772
773        // Return the result.
774        receive_worker.await
775    }
776
777    #[test_case(|stack_trace| stack_trace.stack_trace_key = None => matches
778        Err(Error::MissingField { container: "StackTrace", field: "stack_trace_key" }) ; "stack_trace_key")]
779    #[test_case(|stack_trace| stack_trace.program_addresses = None => matches
780        Err(Error::MissingField { container: "StackTrace", field: "program_addresses" }) ; "program_addresses")]
781    #[test_case(|_| () /* if we do not set any field to None, the result should be Ok */ => matches
782        Ok(_) ; "success")]
783    #[fasync::run_singlethreaded(test)]
784    async fn test_stack_trace_required_fields(
785        set_one_field_to_none: fn(&mut fheapdump_client::StackTrace),
786    ) -> Result<Snapshot, Error> {
787        let client = create_client();
788        let (receiver_proxy, receiver_stream) =
789            client.create_proxy_and_stream::<fheapdump_client::SnapshotReceiverMarker>();
790        let receive_worker = fasync::Task::local(Snapshot::receive_single_from(receiver_stream));
791
792        // Start with a StackTrace with all the required fields set.
793        let mut stack_trace = fheapdump_client::StackTrace {
794            stack_trace_key: Some(FAKE_STACK_TRACE_1_KEY),
795            program_addresses: Some(FAKE_STACK_TRACE_1_ADDRESSES.to_vec()),
796            ..Default::default()
797        };
798
799        // Set one of the fields to None, according to the case being tested.
800        set_one_field_to_none(&mut stack_trace);
801
802        // Send it to the SnapshotReceiver.
803        let fut =
804            receiver_proxy.batch(&[fheapdump_client::SnapshotElement::StackTrace(stack_trace)]);
805        let _ = fut.await; // ignore result, as the peer may detect the error and close the channel
806
807        // Send the end of stream marker.
808        let fut = receiver_proxy.batch(&[]);
809        let _ = fut.await; // ignore result, as the peer may detect the error and close the channel
810
811        // Return the result.
812        receive_worker.await
813    }
814
815    #[test_case(|region| region.address = None => matches
816        Err(Error::MissingField { container: "ExecutableRegion", field: "address" }) ; "address")]
817    #[test_case(|region| region.size = None => matches
818        Err(Error::MissingField { container: "ExecutableRegion", field: "size" }) ; "size")]
819    #[test_case(|region| region.file_offset = None => matches
820        Err(Error::MissingField { container: "ExecutableRegion", field: "file_offset" }) ; "file_offset")]
821    #[test_case(|region| region.build_id = None => matches
822        Err(Error::MissingField { container: "ExecutableRegion", field: "build_id" }) ; "build_id")]
823    #[test_case(|region| region.vaddr = None => matches
824        Err(Error::MissingField { container: "ExecutableRegion", field: "vaddr" }) ; "vaddr")]
825    #[test_case(|_| () /* if we do not set any field to None, the result should be Ok */ => matches
826        Ok(_) ; "success")]
827    #[fasync::run_singlethreaded(test)]
828    async fn test_executable_region_required_fields(
829        set_one_field_to_none: fn(&mut fheapdump_client::ExecutableRegion),
830    ) -> Result<Snapshot, Error> {
831        let client = create_client();
832        let (receiver_proxy, receiver_stream) =
833            client.create_proxy_and_stream::<fheapdump_client::SnapshotReceiverMarker>();
834        let receive_worker = fasync::Task::local(Snapshot::receive_single_from(receiver_stream));
835
836        // Start with an ExecutableRegion with all the required fields set.
837        let mut region = fheapdump_client::ExecutableRegion {
838            address: Some(FAKE_REGION_1_ADDRESS),
839            size: Some(FAKE_REGION_1_SIZE),
840            file_offset: Some(FAKE_REGION_1_FILE_OFFSET),
841            build_id: Some(fheapdump_client::BuildId { value: FAKE_REGION_1_BUILD_ID.to_vec() }),
842            vaddr: Some(FAKE_REGION_1_VADDR),
843            ..Default::default()
844        };
845
846        // Set one of the fields to None, according to the case being tested.
847        set_one_field_to_none(&mut region);
848
849        // Send it to the SnapshotReceiver.
850        let fut =
851            receiver_proxy.batch(&[fheapdump_client::SnapshotElement::ExecutableRegion(region)]);
852        let _ = fut.await; // ignore result, as the peer may detect the error and close the channel
853
854        // Send the end of stream marker.
855        let fut = receiver_proxy.batch(&[]);
856        let _ = fut.await; // ignore result, as the peer may detect the error and close the channel
857
858        // Return the result.
859        receive_worker.await
860    }
861
862    #[test_case(|block_contents| block_contents.address = None => matches
863        Err(Error::MissingField { container: "BlockContents", field: "address" }) ; "address")]
864    #[test_case(|block_contents| block_contents.contents = None => matches
865        Err(Error::MissingField { container: "BlockContents", field: "contents" }) ; "contents")]
866    #[test_case(|_| () /* if we do not set any field to None, the result should be Ok */ => matches
867        Ok(_) ; "success")]
868    #[fasync::run_singlethreaded(test)]
869    async fn test_block_contents_required_fields(
870        set_one_field_to_none: fn(&mut fheapdump_client::BlockContents),
871    ) -> Result<Snapshot, Error> {
872        let client = create_client();
873        let (receiver_proxy, receiver_stream) =
874            client.create_proxy_and_stream::<fheapdump_client::SnapshotReceiverMarker>();
875        let receive_worker = fasync::Task::local(Snapshot::receive_single_from(receiver_stream));
876
877        // Start with a BlockContents with all the required fields set.
878        let mut block_contents = fheapdump_client::BlockContents {
879            address: Some(FAKE_ALLOCATION_1_ADDRESS),
880            contents: Some(FAKE_ALLOCATION_1_CONTENTS.to_vec()),
881            ..Default::default()
882        };
883
884        // Set one of the fields to None, according to the case being tested.
885        set_one_field_to_none(&mut block_contents);
886
887        // Send it to the SnapshotReceiver along with the allocation it references.
888        let fut = receiver_proxy.batch(&[
889            fheapdump_client::SnapshotElement::BlockContents(block_contents),
890            fheapdump_client::SnapshotElement::Allocation(fheapdump_client::Allocation {
891                address: Some(FAKE_ALLOCATION_1_ADDRESS),
892                size: Some(FAKE_ALLOCATION_1_SIZE),
893                thread_info_key: Some(FAKE_THREAD_1_KEY),
894                stack_trace_key: Some(FAKE_STACK_TRACE_1_KEY),
895                timestamp: Some(FAKE_ALLOCATION_1_TIMESTAMP),
896                ..Default::default()
897            }),
898            fheapdump_client::SnapshotElement::ThreadInfo(fheapdump_client::ThreadInfo {
899                thread_info_key: Some(FAKE_THREAD_1_KEY),
900                koid: Some(FAKE_THREAD_1_KOID),
901                name: Some(FAKE_THREAD_1_NAME.to_string()),
902                ..Default::default()
903            }),
904            fheapdump_client::SnapshotElement::StackTrace(fheapdump_client::StackTrace {
905                stack_trace_key: Some(FAKE_STACK_TRACE_1_KEY),
906                program_addresses: Some(FAKE_STACK_TRACE_1_ADDRESSES.to_vec()),
907                ..Default::default()
908            }),
909        ]);
910        let _ = fut.await; // ignore result, as the peer may detect the error and close the channel
911
912        // Send the end of stream marker.
913        let fut = receiver_proxy.batch(&[]);
914        let _ = fut.await; // ignore result, as the peer may detect the error and close the channel
915
916        // Return the result.
917        receive_worker.await
918    }
919
920    #[fasync::run_singlethreaded(test)]
921    async fn test_conflicting_allocations() {
922        let client = create_client();
923        let (receiver_proxy, receiver_stream) =
924            client.create_proxy_and_stream::<fheapdump_client::SnapshotReceiverMarker>();
925        let receive_worker = fasync::Task::local(Snapshot::receive_single_from(receiver_stream));
926
927        // Send two allocations with the same address along with the stack trace they reference.
928        let fut = receiver_proxy.batch(&[
929            fheapdump_client::SnapshotElement::Allocation(fheapdump_client::Allocation {
930                address: Some(FAKE_ALLOCATION_1_ADDRESS),
931                size: Some(FAKE_ALLOCATION_1_SIZE),
932                thread_info_key: Some(FAKE_THREAD_1_KEY),
933                stack_trace_key: Some(FAKE_STACK_TRACE_1_KEY),
934                timestamp: Some(FAKE_ALLOCATION_1_TIMESTAMP),
935                ..Default::default()
936            }),
937            fheapdump_client::SnapshotElement::Allocation(fheapdump_client::Allocation {
938                address: Some(FAKE_ALLOCATION_1_ADDRESS),
939                size: Some(FAKE_ALLOCATION_1_SIZE),
940                thread_info_key: Some(FAKE_THREAD_1_KEY),
941                stack_trace_key: Some(FAKE_STACK_TRACE_1_KEY),
942                timestamp: Some(FAKE_ALLOCATION_1_TIMESTAMP),
943                ..Default::default()
944            }),
945            fheapdump_client::SnapshotElement::ThreadInfo(fheapdump_client::ThreadInfo {
946                thread_info_key: Some(FAKE_THREAD_1_KEY),
947                koid: Some(FAKE_THREAD_1_KOID),
948                name: Some(FAKE_THREAD_1_NAME.to_string()),
949                ..Default::default()
950            }),
951            fheapdump_client::SnapshotElement::StackTrace(fheapdump_client::StackTrace {
952                stack_trace_key: Some(FAKE_STACK_TRACE_1_KEY),
953                program_addresses: Some(FAKE_STACK_TRACE_1_ADDRESSES.to_vec()),
954                ..Default::default()
955            }),
956        ]);
957        let _ = fut.await; // ignore result, as the peer may detect the error and close the channel
958
959        // Send the end of stream marker.
960        let fut = receiver_proxy.batch(&[]);
961        let _ = fut.await; // ignore result, as the peer may detect the error and close the channel
962
963        // Verify expected error.
964        assert_matches!(
965            receive_worker.await,
966            Err(Error::ConflictingElement { element_type: "Allocation" })
967        );
968    }
969
970    #[fasync::run_singlethreaded(test)]
971    async fn test_conflicting_executable_regions() {
972        let client = create_client();
973        let (receiver_proxy, receiver_stream) =
974            client.create_proxy_and_stream::<fheapdump_client::SnapshotReceiverMarker>();
975        let receive_worker = fasync::Task::local(Snapshot::receive_single_from(receiver_stream));
976
977        // Send two executable regions with the same address.
978        let fut = receiver_proxy.batch(&[
979            fheapdump_client::SnapshotElement::ExecutableRegion(
980                fheapdump_client::ExecutableRegion {
981                    address: Some(FAKE_REGION_1_ADDRESS),
982                    size: Some(FAKE_REGION_1_SIZE),
983                    file_offset: Some(FAKE_REGION_1_FILE_OFFSET),
984                    build_id: Some(fheapdump_client::BuildId {
985                        value: FAKE_REGION_1_BUILD_ID.to_vec(),
986                    }),
987                    vaddr: Some(FAKE_REGION_1_VADDR),
988                    ..Default::default()
989                },
990            ),
991            fheapdump_client::SnapshotElement::ExecutableRegion(
992                fheapdump_client::ExecutableRegion {
993                    address: Some(FAKE_REGION_1_ADDRESS),
994                    size: Some(FAKE_REGION_1_SIZE),
995                    file_offset: Some(FAKE_REGION_1_FILE_OFFSET),
996                    build_id: Some(fheapdump_client::BuildId {
997                        value: FAKE_REGION_1_BUILD_ID.to_vec(),
998                    }),
999                    vaddr: Some(FAKE_REGION_1_VADDR),
1000                    ..Default::default()
1001                },
1002            ),
1003        ]);
1004        let _ = fut.await; // ignore result, as the peer may detect the error and close the channel
1005
1006        // Send the end of stream marker.
1007        let fut = receiver_proxy.batch(&[]);
1008        let _ = fut.await; // ignore result, as the peer may detect the error and close the channel
1009
1010        // Verify expected error.
1011        assert_matches!(
1012            receive_worker.await,
1013            Err(Error::ConflictingElement { element_type: "ExecutableRegion" })
1014        );
1015    }
1016
1017    #[fasync::run_singlethreaded(test)]
1018    async fn test_block_contents_wrong_size() {
1019        let client = create_client();
1020        let (receiver_proxy, receiver_stream) =
1021            client.create_proxy_and_stream::<fheapdump_client::SnapshotReceiverMarker>();
1022        let receive_worker = fasync::Task::local(Snapshot::receive_single_from(receiver_stream));
1023
1024        // Send an allocation whose BlockContents has the wrong size.
1025        let contents_with_wrong_size = vec![0; FAKE_ALLOCATION_1_SIZE as usize + 1];
1026        let fut = receiver_proxy.batch(&[
1027            fheapdump_client::SnapshotElement::Allocation(fheapdump_client::Allocation {
1028                address: Some(FAKE_ALLOCATION_1_ADDRESS),
1029                size: Some(FAKE_ALLOCATION_1_SIZE),
1030                thread_info_key: Some(FAKE_THREAD_1_KEY),
1031                stack_trace_key: Some(FAKE_STACK_TRACE_1_KEY),
1032                timestamp: Some(FAKE_ALLOCATION_1_TIMESTAMP),
1033                ..Default::default()
1034            }),
1035            fheapdump_client::SnapshotElement::BlockContents(fheapdump_client::BlockContents {
1036                address: Some(FAKE_ALLOCATION_1_ADDRESS),
1037                contents: Some(contents_with_wrong_size),
1038                ..Default::default()
1039            }),
1040            fheapdump_client::SnapshotElement::ThreadInfo(fheapdump_client::ThreadInfo {
1041                thread_info_key: Some(FAKE_THREAD_1_KEY),
1042                koid: Some(FAKE_THREAD_1_KOID),
1043                name: Some(FAKE_THREAD_1_NAME.to_string()),
1044                ..Default::default()
1045            }),
1046            fheapdump_client::SnapshotElement::StackTrace(fheapdump_client::StackTrace {
1047                stack_trace_key: Some(FAKE_STACK_TRACE_1_KEY),
1048                program_addresses: Some(FAKE_STACK_TRACE_1_ADDRESSES.to_vec()),
1049                ..Default::default()
1050            }),
1051        ]);
1052        let _ = fut.await; // ignore result, as the peer may detect the error and close the channel
1053
1054        // Send the end of stream marker.
1055        let fut = receiver_proxy.batch(&[]);
1056        let _ = fut.await; // ignore result, as the peer may detect the error and close the channel
1057
1058        // Verify expected error.
1059        assert_matches!(
1060            receive_worker.await,
1061            Err(Error::ConflictingElement { element_type: "BlockContents" })
1062        );
1063    }
1064
1065    #[fasync::run_singlethreaded(test)]
1066    async fn test_empty_stack_trace() {
1067        let client = create_client();
1068        let (receiver_proxy, receiver_stream) =
1069            client.create_proxy_and_stream::<fheapdump_client::SnapshotReceiverMarker>();
1070        let receive_worker = fasync::Task::local(Snapshot::receive_single_from(receiver_stream));
1071
1072        // Send an allocation that references an empty stack trace.
1073        let fut = receiver_proxy.batch(&[
1074            fheapdump_client::SnapshotElement::Allocation(fheapdump_client::Allocation {
1075                address: Some(FAKE_ALLOCATION_1_ADDRESS),
1076                size: Some(FAKE_ALLOCATION_1_SIZE),
1077                thread_info_key: Some(FAKE_THREAD_1_KEY),
1078                stack_trace_key: Some(FAKE_STACK_TRACE_1_KEY),
1079                timestamp: Some(FAKE_ALLOCATION_1_TIMESTAMP),
1080                ..Default::default()
1081            }),
1082            fheapdump_client::SnapshotElement::ThreadInfo(fheapdump_client::ThreadInfo {
1083                thread_info_key: Some(FAKE_THREAD_1_KEY),
1084                koid: Some(FAKE_THREAD_1_KOID),
1085                name: Some(FAKE_THREAD_1_NAME.to_string()),
1086                ..Default::default()
1087            }),
1088            fheapdump_client::SnapshotElement::StackTrace(fheapdump_client::StackTrace {
1089                stack_trace_key: Some(FAKE_STACK_TRACE_1_KEY),
1090                program_addresses: Some(vec![]),
1091                ..Default::default()
1092            }),
1093        ]);
1094        fut.await.unwrap();
1095
1096        // Send the end of stream marker.
1097        let fut = receiver_proxy.batch(&[]);
1098        fut.await.unwrap();
1099
1100        // Verify that the stack trace has been reconstructed correctly.
1101        let received_snapshot = receive_worker.await.unwrap();
1102        let allocation1 = received_snapshot
1103            .allocations
1104            .iter()
1105            .find(|a| a.address == Some(FAKE_ALLOCATION_1_ADDRESS))
1106            .unwrap();
1107        assert_eq!(allocation1.stack_trace.program_addresses, []);
1108    }
1109
1110    #[fasync::run_singlethreaded(test)]
1111    async fn test_chunked_stack_trace() {
1112        let client = create_client();
1113        let (receiver_proxy, receiver_stream) =
1114            client.create_proxy_and_stream::<fheapdump_client::SnapshotReceiverMarker>();
1115        let receive_worker = fasync::Task::local(Snapshot::receive_single_from(receiver_stream));
1116
1117        // Send an allocation and the first chunk of its stack trace.
1118        let fut = receiver_proxy.batch(&[
1119            fheapdump_client::SnapshotElement::Allocation(fheapdump_client::Allocation {
1120                address: Some(FAKE_ALLOCATION_1_ADDRESS),
1121                size: Some(FAKE_ALLOCATION_1_SIZE),
1122                thread_info_key: Some(FAKE_THREAD_1_KEY),
1123                stack_trace_key: Some(FAKE_STACK_TRACE_1_KEY),
1124                timestamp: Some(FAKE_ALLOCATION_1_TIMESTAMP),
1125                ..Default::default()
1126            }),
1127            fheapdump_client::SnapshotElement::StackTrace(fheapdump_client::StackTrace {
1128                stack_trace_key: Some(FAKE_STACK_TRACE_1_KEY),
1129                program_addresses: Some(vec![1111, 2222]),
1130                ..Default::default()
1131            }),
1132            fheapdump_client::SnapshotElement::ThreadInfo(fheapdump_client::ThreadInfo {
1133                thread_info_key: Some(FAKE_THREAD_1_KEY),
1134                koid: Some(FAKE_THREAD_1_KOID),
1135                name: Some(FAKE_THREAD_1_NAME.to_string()),
1136                ..Default::default()
1137            }),
1138        ]);
1139        fut.await.unwrap();
1140
1141        // Send the second chunk.
1142        let fut = receiver_proxy.batch(&[fheapdump_client::SnapshotElement::StackTrace(
1143            fheapdump_client::StackTrace {
1144                stack_trace_key: Some(FAKE_STACK_TRACE_1_KEY),
1145                program_addresses: Some(vec![3333]),
1146                ..Default::default()
1147            },
1148        )]);
1149        fut.await.unwrap();
1150
1151        // Send the end of stream marker.
1152        let fut = receiver_proxy.batch(&[]);
1153        fut.await.unwrap();
1154
1155        // Verify that the stack trace has been reconstructed correctly.
1156        let received_snapshot = receive_worker.await.unwrap();
1157        let allocation1 = received_snapshot
1158            .allocations
1159            .iter()
1160            .find(|alloc| alloc.address == Some(FAKE_ALLOCATION_1_ADDRESS))
1161            .unwrap();
1162        assert_eq!(allocation1.stack_trace.program_addresses, [1111, 2222, 3333]);
1163    }
1164
1165    #[fasync::run_singlethreaded(test)]
1166    async fn test_empty_block_contents() {
1167        let client = create_client();
1168        let (receiver_proxy, receiver_stream) =
1169            client.create_proxy_and_stream::<fheapdump_client::SnapshotReceiverMarker>();
1170        let receive_worker = fasync::Task::local(Snapshot::receive_single_from(receiver_stream));
1171
1172        // Send a zero-sized allocation and its empty contents.
1173        let fut = receiver_proxy.batch(&[
1174            fheapdump_client::SnapshotElement::Allocation(fheapdump_client::Allocation {
1175                address: Some(FAKE_ALLOCATION_1_ADDRESS),
1176                size: Some(0),
1177                thread_info_key: Some(FAKE_THREAD_1_KEY),
1178                stack_trace_key: Some(FAKE_STACK_TRACE_1_KEY),
1179                timestamp: Some(FAKE_ALLOCATION_1_TIMESTAMP),
1180                ..Default::default()
1181            }),
1182            fheapdump_client::SnapshotElement::ThreadInfo(fheapdump_client::ThreadInfo {
1183                thread_info_key: Some(FAKE_THREAD_1_KEY),
1184                koid: Some(FAKE_THREAD_1_KOID),
1185                name: Some(FAKE_THREAD_1_NAME.to_string()),
1186                ..Default::default()
1187            }),
1188            fheapdump_client::SnapshotElement::StackTrace(fheapdump_client::StackTrace {
1189                stack_trace_key: Some(FAKE_STACK_TRACE_1_KEY),
1190                program_addresses: Some(FAKE_STACK_TRACE_1_ADDRESSES.to_vec()),
1191                ..Default::default()
1192            }),
1193            fheapdump_client::SnapshotElement::BlockContents(fheapdump_client::BlockContents {
1194                address: Some(FAKE_ALLOCATION_1_ADDRESS),
1195                contents: Some(vec![]),
1196                ..Default::default()
1197            }),
1198        ]);
1199        fut.await.unwrap();
1200
1201        // Send the end of stream marker.
1202        let fut = receiver_proxy.batch(&[]);
1203        fut.await.unwrap();
1204
1205        // Verify that the allocation has been reconstructed correctly.
1206        let received_snapshot = receive_worker.await.unwrap();
1207        let allocation1 = received_snapshot
1208            .allocations
1209            .iter()
1210            .find(|alloc| alloc.address == Some(FAKE_ALLOCATION_1_ADDRESS))
1211            .unwrap();
1212        assert_eq!(allocation1.contents.as_ref().expect("contents must be set"), &vec![]);
1213    }
1214
1215    #[fasync::run_singlethreaded(test)]
1216    async fn test_chunked_block_contents() {
1217        let client = create_client();
1218        let (receiver_proxy, receiver_stream) =
1219            client.create_proxy_and_stream::<fheapdump_client::SnapshotReceiverMarker>();
1220        let receive_worker = fasync::Task::local(Snapshot::receive_single_from(receiver_stream));
1221
1222        // Split the contents in two halves.
1223        let (content_first_chunk, contents_second_chunk) =
1224            FAKE_ALLOCATION_1_CONTENTS.split_at(FAKE_ALLOCATION_1_CONTENTS.len() / 2);
1225
1226        // Send an allocation and the first chunk of its contents.
1227        let fut = receiver_proxy.batch(&[
1228            fheapdump_client::SnapshotElement::Allocation(fheapdump_client::Allocation {
1229                address: Some(FAKE_ALLOCATION_1_ADDRESS),
1230                size: Some(FAKE_ALLOCATION_1_SIZE),
1231                thread_info_key: Some(FAKE_THREAD_1_KEY),
1232                stack_trace_key: Some(FAKE_STACK_TRACE_1_KEY),
1233                timestamp: Some(FAKE_ALLOCATION_1_TIMESTAMP),
1234                ..Default::default()
1235            }),
1236            fheapdump_client::SnapshotElement::ThreadInfo(fheapdump_client::ThreadInfo {
1237                thread_info_key: Some(FAKE_THREAD_1_KEY),
1238                koid: Some(FAKE_THREAD_1_KOID),
1239                name: Some(FAKE_THREAD_1_NAME.to_string()),
1240                ..Default::default()
1241            }),
1242            fheapdump_client::SnapshotElement::StackTrace(fheapdump_client::StackTrace {
1243                stack_trace_key: Some(FAKE_STACK_TRACE_1_KEY),
1244                program_addresses: Some(FAKE_STACK_TRACE_1_ADDRESSES.to_vec()),
1245                ..Default::default()
1246            }),
1247            fheapdump_client::SnapshotElement::BlockContents(fheapdump_client::BlockContents {
1248                address: Some(FAKE_ALLOCATION_1_ADDRESS),
1249                contents: Some(content_first_chunk.to_vec()),
1250                ..Default::default()
1251            }),
1252        ]);
1253        fut.await.unwrap();
1254
1255        // Send the second chunk.
1256        let fut = receiver_proxy.batch(&[fheapdump_client::SnapshotElement::BlockContents(
1257            fheapdump_client::BlockContents {
1258                address: Some(FAKE_ALLOCATION_1_ADDRESS),
1259                contents: Some(contents_second_chunk.to_vec()),
1260                ..Default::default()
1261            },
1262        )]);
1263        fut.await.unwrap();
1264
1265        // Send the end of stream marker.
1266        let fut = receiver_proxy.batch(&[]);
1267        fut.await.unwrap();
1268
1269        // Verify that the allocation's block contents have been reconstructed correctly.
1270        let received_snapshot = receive_worker.await.unwrap();
1271        let allocation1 = received_snapshot
1272            .allocations
1273            .iter()
1274            .find(|a| a.address == Some(FAKE_ALLOCATION_1_ADDRESS))
1275            .unwrap();
1276        assert_eq!(allocation1.contents, Some(FAKE_ALLOCATION_1_CONTENTS.to_vec()));
1277    }
1278
1279    #[fasync::run_singlethreaded(test)]
1280    async fn test_missing_end_of_stream() {
1281        let client = create_client();
1282        let (receiver_proxy, receiver_stream) =
1283            client.create_proxy_and_stream::<fheapdump_client::SnapshotReceiverMarker>();
1284        let receive_worker = fasync::Task::local(Snapshot::receive_single_from(receiver_stream));
1285
1286        // Send an allocation and its stack trace.
1287        let fut = receiver_proxy.batch(&[
1288            fheapdump_client::SnapshotElement::Allocation(fheapdump_client::Allocation {
1289                address: Some(FAKE_ALLOCATION_1_ADDRESS),
1290                size: Some(FAKE_ALLOCATION_1_SIZE),
1291                thread_info_key: Some(FAKE_THREAD_1_KEY),
1292                stack_trace_key: Some(FAKE_STACK_TRACE_1_KEY),
1293                timestamp: Some(FAKE_ALLOCATION_1_TIMESTAMP),
1294                ..Default::default()
1295            }),
1296            fheapdump_client::SnapshotElement::ThreadInfo(fheapdump_client::ThreadInfo {
1297                thread_info_key: Some(FAKE_THREAD_1_KEY),
1298                koid: Some(FAKE_THREAD_1_KOID),
1299                name: Some(FAKE_THREAD_1_NAME.to_string()),
1300                ..Default::default()
1301            }),
1302            fheapdump_client::SnapshotElement::StackTrace(fheapdump_client::StackTrace {
1303                stack_trace_key: Some(FAKE_STACK_TRACE_1_KEY),
1304                program_addresses: Some(FAKE_STACK_TRACE_1_ADDRESSES.to_vec()),
1305                ..Default::default()
1306            }),
1307        ]);
1308        fut.await.unwrap();
1309
1310        // Close the channel without sending an end of stream marker.
1311        std::mem::drop(receiver_proxy);
1312
1313        // Expect an UnexpectedEndOfStream error.
1314        assert_matches!(receive_worker.await, Err(Error::UnexpectedEndOfStream));
1315    }
1316
1317    #[fasync::run_singlethreaded(test)]
1318    async fn test_multi_contents() {
1319        let client = create_client();
1320        let (receiver_proxy, receiver_stream) =
1321            client.create_proxy_and_stream::<fheapdump_client::SnapshotReceiverMarker>();
1322        let receive_worker = fasync::Task::local(Snapshot::receive_multi_from(receiver_stream));
1323
1324        // Send two snapshots with different KOIDs.
1325        for koid in [1111, 2222] {
1326            receiver_proxy
1327                .batch(&[fheapdump_client::SnapshotElement::SnapshotHeader(
1328                    fheapdump_client::SnapshotHeader {
1329                        process_name: Some(format!("test-process-{koid}")),
1330                        process_koid: Some(koid),
1331                        ..Default::default()
1332                    },
1333                )])
1334                .await
1335                .unwrap();
1336
1337            receiver_proxy
1338                .batch(&[
1339                    fheapdump_client::SnapshotElement::Allocation(fheapdump_client::Allocation {
1340                        address: Some(FAKE_ALLOCATION_1_ADDRESS),
1341                        size: Some(FAKE_ALLOCATION_1_SIZE),
1342                        thread_info_key: Some(FAKE_THREAD_1_KEY),
1343                        stack_trace_key: Some(FAKE_STACK_TRACE_1_KEY),
1344                        timestamp: Some(FAKE_ALLOCATION_1_TIMESTAMP),
1345                        ..Default::default()
1346                    }),
1347                    fheapdump_client::SnapshotElement::ThreadInfo(fheapdump_client::ThreadInfo {
1348                        thread_info_key: Some(FAKE_THREAD_1_KEY),
1349                        koid: Some(FAKE_THREAD_1_KOID),
1350                        name: Some(FAKE_THREAD_1_NAME.to_string()),
1351                        ..Default::default()
1352                    }),
1353                    fheapdump_client::SnapshotElement::StackTrace(fheapdump_client::StackTrace {
1354                        stack_trace_key: Some(FAKE_STACK_TRACE_1_KEY),
1355                        program_addresses: Some(FAKE_STACK_TRACE_1_ADDRESSES.to_vec()),
1356                        ..Default::default()
1357                    }),
1358                ])
1359                .await
1360                .unwrap();
1361
1362            receiver_proxy.batch(&[]).await.unwrap(); // end of snapshot
1363        }
1364
1365        // Send end of stream marker.
1366        receiver_proxy.batch(&[]).await.unwrap();
1367
1368        // Validate the received snapshots.
1369        let received_snapshots = receive_worker.await.unwrap();
1370        assert_eq!(received_snapshots.len(), 2);
1371        assert_eq!(received_snapshots[0].process_name, "test-process-1111");
1372        assert_eq!(received_snapshots[0].process_koid, 1111);
1373        assert_eq!(received_snapshots[1].process_name, "test-process-2222");
1374        assert_eq!(received_snapshots[1].process_koid, 2222);
1375    }
1376
1377    #[fasync::run_singlethreaded(test)]
1378    async fn test_multi_missing_end_of_stream() {
1379        let client = create_client();
1380        let (receiver_proxy, receiver_stream) =
1381            client.create_proxy_and_stream::<fheapdump_client::SnapshotReceiverMarker>();
1382        let receive_worker = fasync::Task::local(Snapshot::receive_multi_from(receiver_stream));
1383
1384        // Send two snapshots with different KOIDs.
1385        for koid in [1111, 2222] {
1386            receiver_proxy
1387                .batch(&[fheapdump_client::SnapshotElement::SnapshotHeader(
1388                    fheapdump_client::SnapshotHeader {
1389                        process_name: Some("test-process-name".to_string()),
1390                        process_koid: Some(koid),
1391                        ..Default::default()
1392                    },
1393                )])
1394                .await
1395                .unwrap();
1396
1397            receiver_proxy
1398                .batch(&[
1399                    fheapdump_client::SnapshotElement::Allocation(fheapdump_client::Allocation {
1400                        address: Some(FAKE_ALLOCATION_1_ADDRESS),
1401                        size: Some(FAKE_ALLOCATION_1_SIZE),
1402                        thread_info_key: Some(FAKE_THREAD_1_KEY),
1403                        stack_trace_key: Some(FAKE_STACK_TRACE_1_KEY),
1404                        timestamp: Some(FAKE_ALLOCATION_1_TIMESTAMP),
1405                        ..Default::default()
1406                    }),
1407                    fheapdump_client::SnapshotElement::ThreadInfo(fheapdump_client::ThreadInfo {
1408                        thread_info_key: Some(FAKE_THREAD_1_KEY),
1409                        koid: Some(FAKE_THREAD_1_KOID),
1410                        name: Some(FAKE_THREAD_1_NAME.to_string()),
1411                        ..Default::default()
1412                    }),
1413                    fheapdump_client::SnapshotElement::StackTrace(fheapdump_client::StackTrace {
1414                        stack_trace_key: Some(FAKE_STACK_TRACE_1_KEY),
1415                        program_addresses: Some(FAKE_STACK_TRACE_1_ADDRESSES.to_vec()),
1416                        ..Default::default()
1417                    }),
1418                ])
1419                .await
1420                .unwrap();
1421
1422            receiver_proxy.batch(&[]).await.unwrap(); // end of snapshot
1423        }
1424
1425        // Close the channel without sending an end of stream marker.
1426        std::mem::drop(receiver_proxy);
1427
1428        // Expect an UnexpectedEndOfStream error.
1429        assert_matches!(receive_worker.await, Err(Error::UnexpectedEndOfStream));
1430    }
1431}