Skip to main content

perfetto/
lib.rs

1// Copyright 2024 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
5// Increase recursion limit because LTO causes overflow.
6#![recursion_limit = "256"]
7
8use perfetto_protos::perfetto::protos::{
9    DisableTracingRequest, EnableTracingRequest, FreeBuffersRequest, GetAsyncCommandRequest,
10    GetAsyncCommandResponse, InitializeConnectionRequest, InitializeConnectionResponse, IpcFrame,
11    ReadBuffersRequest, RegisterDataSourceRequest, RegisterDataSourceResponse, ipc_frame,
12};
13use prost::Message;
14use starnix_core::task::{CurrentTask, EventHandler, Waiter};
15use starnix_core::vfs::buffers::{VecInputBuffer, VecOutputBuffer};
16use starnix_core::vfs::socket::{
17    SocketDomain, SocketFile, SocketPeer, SocketProtocol, SocketType, resolve_unix_socket_address,
18};
19use starnix_core::vfs::{FileHandle, FsStr};
20
21use starnix_uapi::errors::Errno;
22use starnix_uapi::open_flags::OpenFlags;
23use starnix_uapi::vfs::FdEvents;
24use std::collections::VecDeque;
25use thiserror::Error;
26
27/// A connection represents serializing Perfetto IPC requests to a file
28/// Specifically, this is writing 4-byte-little-endian length prefixed proto-coded messages to the
29/// receiver.
30///
31/// See https://perfetto.dev/docs/design-docs/api-and-abi for more details
32pub struct IpcConnection {
33    /// File to write to.
34    file: FileHandle,
35    /// Next unused request id. This is used for correlating replies to requests. We need to
36    /// increment this after each request.
37    request_id: u64,
38}
39
40#[derive(Error, Debug)]
41pub enum IpcWriteError {
42    #[error(transparent)]
43    Encode(#[from] prost::EncodeError),
44    #[error(transparent)]
45    Write(#[from] Errno),
46    #[error("TooLong: {0} exceeds max for u32")]
47    TooLong(usize),
48}
49
50#[derive(Error, Debug)]
51pub enum IpcReadError {
52    #[error(transparent)]
53    Decode(#[from] prost::DecodeError),
54    #[error(transparent)]
55    Read(#[from] Errno),
56}
57
58#[derive(Error, Debug)]
59pub enum InvokeMethodError {
60    #[error("could not not look up method name: {0}")]
61    InvalidMethod(String),
62    #[error(transparent)]
63    IpcWrite(#[from] IpcWriteError),
64    #[error(transparent)]
65    IpcRead(#[from] IpcReadError),
66    #[error("unexpected response: {0}")]
67    InvalidResponse(String),
68}
69
70#[derive(Error, Debug)]
71pub enum ProducerError {
72    #[error(transparent)]
73    InvokeMethod(#[from] InvokeMethodError),
74    #[error(transparent)]
75    IpcWrite(#[from] IpcWriteError),
76    #[error(transparent)]
77    IpcRead(#[from] IpcReadError),
78    #[error("unexpected response: {0}")]
79    InvalidResponse(String),
80    #[error(transparent)]
81    Decode(#[from] prost::DecodeError),
82}
83
84impl IpcConnection {
85    pub fn new(file: FileHandle) -> Self {
86        Self { file, request_id: 0 }
87    }
88
89    pub fn bind_service(
90        &mut self,
91        service_name: &str,
92        current_task: &CurrentTask,
93    ) -> Result<(), IpcWriteError> {
94        // The first thing we need to send over our newly connected socket is send a request to
95        // bind to the service.
96        let bind_service_message = IpcFrame {
97            request_id: Some(1),
98            msg: Some(ipc_frame::Msg::MsgBindService(ipc_frame::BindService {
99                service_name: Some(service_name.to_string()),
100            })),
101            ..Default::default()
102        };
103        self.write_frame(bind_service_message, current_task)
104    }
105
106    pub fn invoke_method(
107        &mut self,
108        service_id: u32,
109        method_id: u32,
110        arguments: Option<Vec<u8>>,
111        current_task: &CurrentTask,
112    ) -> Result<(), IpcWriteError> {
113        let msg = IpcFrame {
114            request_id: Some(self.allocate_request_id()),
115            msg: Some(ipc_frame::Msg::MsgInvokeMethod(ipc_frame::InvokeMethod {
116                service_id: Some(service_id),
117                method_id: Some(method_id),
118                args_proto: arguments,
119                drop_reply: None,
120            })),
121            ..Default::default()
122        };
123        self.write_frame(msg, current_task)
124    }
125
126    fn write_frame(
127        &mut self,
128        frame: IpcFrame,
129        current_task: &CurrentTask,
130    ) -> Result<(), IpcWriteError> {
131        // We need to length prefix our proto before encoding and sending it down the wire.
132        let frame_len = u32::try_from(frame.encoded_len())
133            .map_err(|_| IpcWriteError::TooLong(frame.encoded_len()))?;
134        let mut bind_service_bytes =
135            Vec::with_capacity(frame.encoded_len() + std::mem::size_of::<u32>());
136        bind_service_bytes.extend_from_slice(&frame_len.to_le_bytes());
137        frame.encode(&mut bind_service_bytes)?;
138        let mut bind_service_buffer: VecInputBuffer = bind_service_bytes.into();
139        self.file.write(current_task, &mut bind_service_buffer)?;
140        Ok(())
141    }
142
143    // Perfetto requires a unique message id for each ipc request.
144    fn allocate_request_id(&mut self) -> u64 {
145        let id = self.request_id;
146        self.request_id += 1;
147        id
148    }
149}
150
151/// State for reading Perfetto IPC frames.
152///
153/// Each frame is composed of a 32 bit length in little endian, followed by
154/// the proto-encoded message. This state handles reads that only include
155/// partial messages.
156pub struct FrameReader {
157    /// File to read from.
158    file: FileHandle,
159    /// Buffer for passing to read() calls.
160    ///
161    /// This buffer does not store any data long-term, but is persisted to
162    /// avoid reallocating the buffer repeatedly.
163    read_buffer: VecOutputBuffer,
164    /// Data that has been read but not processed.
165    data: VecDeque<u8>,
166    /// If we've received enough bytes to know the next message's size, those
167    /// bytes are removed from [data] and the size is populated here.
168    next_message_size: Option<usize>,
169}
170
171impl FrameReader {
172    pub fn new(file: FileHandle) -> Self {
173        Self {
174            file,
175            read_buffer: VecOutputBuffer::new(4096),
176            data: VecDeque::with_capacity(4096),
177            next_message_size: None,
178        }
179    }
180
181    /// Repeatedly reads from the specified file until a full message is available.
182    pub fn next_frame_blocking(
183        &mut self,
184        current_task: &CurrentTask,
185    ) -> Result<IpcFrame, IpcReadError> {
186        loop {
187            if self.next_message_size.is_none() && self.data.len() >= 4 {
188                let len_bytes: [u8; 4] = self
189                    .data
190                    .drain(..4)
191                    .collect::<Vec<_>>()
192                    .try_into()
193                    .expect("self.data has at least 4 elements");
194                self.next_message_size = Some(u32::from_le_bytes(len_bytes) as usize);
195            }
196            if let Some(message_size) = self.next_message_size {
197                if self.data.len() >= message_size {
198                    let message: Vec<u8> = self.data.drain(..message_size).collect();
199                    self.next_message_size = None;
200                    return Ok(IpcFrame::decode(message.as_slice())?);
201                }
202            }
203
204            let waiter = Waiter::new();
205            self.file.wait_async(current_task, &waiter, FdEvents::POLLIN, EventHandler::None);
206            while self.file.query_events(current_task)? & FdEvents::POLLIN != FdEvents::POLLIN {
207                waiter.wait(current_task)?;
208            }
209            self.file.read(current_task, &mut self.read_buffer)?;
210            self.data.extend(self.read_buffer.data());
211            self.read_buffer.reset();
212        }
213    }
214}
215
216/// Bookkeeping information needed for IPC messages to and from Perfetto.
217pub struct Consumer {
218    /// File handle corresponding to the communication socket. Data is written to and read from
219    /// this file.
220    conn_file: FileHandle,
221    /// State for combining read byte data into IPC frames.
222    frame_reader: FrameReader,
223    /// Reply from the BindService call that was made when the connection was opened.
224    /// This call includes ids for the various IPCs that the Perfetto service supports.
225    bind_service_reply: ipc_frame::BindServiceReply,
226    /// Next unused request id. This is used for correlating repies to requests.
227    request_id: u64,
228}
229
230impl Consumer {
231    /// Opens a socket connection to the specified socket path and initializes the requisite
232    /// bookkeeping information.
233    pub fn new(current_task: &CurrentTask, socket_path: &FsStr) -> Result<Self, anyhow::Error> {
234        let conn_file = SocketFile::new_socket(
235            current_task,
236            SocketDomain::Unix,
237            SocketType::Stream,
238            OpenFlags::RDWR,
239            SocketProtocol::from_raw(0),
240            /* kernel_private=*/ false,
241        )?;
242        let conn = SocketFile::get_from_file(&conn_file)?;
243        let peer = SocketPeer::Handle(resolve_unix_socket_address(current_task, socket_path)?);
244        conn.connect(current_task, peer)?;
245        let mut frame_reader = FrameReader::new(conn_file.clone());
246        let mut request_id = 1;
247
248        let bind_service_message = IpcFrame {
249            request_id: Some(request_id),
250            data_for_testing: Vec::new(),
251            msg: Some(ipc_frame::Msg::MsgBindService(ipc_frame::BindService {
252                service_name: Some("ConsumerPort".to_string()),
253            })),
254        };
255        request_id += 1;
256        let mut bind_service_bytes =
257            Vec::with_capacity(bind_service_message.encoded_len() + std::mem::size_of::<u32>());
258        bind_service_bytes.extend_from_slice(
259            &u32::try_from(bind_service_message.encoded_len()).unwrap().to_le_bytes(),
260        );
261        bind_service_message.encode(&mut bind_service_bytes)?;
262        let mut bind_service_buffer: VecInputBuffer = bind_service_bytes.into();
263        conn.file().write(current_task, &mut bind_service_buffer)?;
264
265        let reply_frame = frame_reader.next_frame_blocking(current_task)?;
266
267        let bind_service_reply = match reply_frame.msg {
268            Some(ipc_frame::Msg::MsgBindServiceReply(reply)) => reply,
269            m => return Err(anyhow::anyhow!("Got unexpected reply message: {:?}", m)),
270        };
271
272        Ok(Self { conn_file, frame_reader, bind_service_reply, request_id })
273    }
274
275    fn send_message(
276        &mut self,
277        current_task: &CurrentTask,
278        msg: ipc_frame::Msg,
279    ) -> Result<u64, anyhow::Error> {
280        let request_id = self.request_id;
281        let frame =
282            IpcFrame { request_id: Some(request_id), data_for_testing: Vec::new(), msg: Some(msg) };
283
284        self.request_id += 1;
285
286        let mut frame_bytes = Vec::with_capacity(frame.encoded_len() + std::mem::size_of::<u32>());
287        frame_bytes.extend_from_slice(&u32::try_from(frame.encoded_len())?.to_le_bytes());
288        frame.encode(&mut frame_bytes)?;
289        let mut buffer: VecInputBuffer = frame_bytes.into();
290        self.conn_file.write(current_task, &mut buffer)?;
291
292        Ok(request_id)
293    }
294
295    fn method_id(&self, name: &str) -> Result<u32, anyhow::Error> {
296        for method in &self.bind_service_reply.methods {
297            if let Some(method_name) = method.name.as_ref() {
298                if method_name == name {
299                    if let Some(id) = method.id {
300                        return Ok(id);
301                    } else {
302                        return Err(anyhow::anyhow!(
303                            "Matched method name {} but found no id",
304                            method_name
305                        ));
306                    }
307                }
308            }
309        }
310        Err(anyhow::anyhow!("Did not find method {}", name))
311    }
312
313    pub fn enable_tracing(
314        &mut self,
315        current_task: &CurrentTask,
316        req: EnableTracingRequest,
317    ) -> Result<u64, anyhow::Error> {
318        let method_id = self.method_id("EnableTracing")?;
319        let mut encoded_args: Vec<u8> = Vec::with_capacity(req.encoded_len());
320        req.encode(&mut encoded_args)?;
321
322        self.send_message(
323            current_task,
324            ipc_frame::Msg::MsgInvokeMethod(ipc_frame::InvokeMethod {
325                service_id: self.bind_service_reply.service_id,
326                method_id: Some(method_id),
327                args_proto: Some(encoded_args),
328                drop_reply: None,
329            }),
330        )
331    }
332
333    pub fn disable_tracing(
334        &mut self,
335        current_task: &CurrentTask,
336        req: DisableTracingRequest,
337    ) -> Result<u64, anyhow::Error> {
338        let method_id = self.method_id("DisableTracing")?;
339        let mut encoded_args: Vec<u8> = Vec::with_capacity(req.encoded_len());
340        req.encode(&mut encoded_args)?;
341
342        self.send_message(
343            current_task,
344            ipc_frame::Msg::MsgInvokeMethod(ipc_frame::InvokeMethod {
345                service_id: self.bind_service_reply.service_id,
346                method_id: Some(method_id),
347                args_proto: Some(encoded_args),
348                drop_reply: None,
349            }),
350        )
351    }
352
353    pub fn read_buffers(
354        &mut self,
355        current_task: &CurrentTask,
356        req: ReadBuffersRequest,
357    ) -> Result<u64, anyhow::Error> {
358        let method_id = self.method_id("ReadBuffers")?;
359        let mut encoded_args: Vec<u8> = Vec::with_capacity(req.encoded_len());
360        req.encode(&mut encoded_args)?;
361
362        self.send_message(
363            current_task,
364            ipc_frame::Msg::MsgInvokeMethod(ipc_frame::InvokeMethod {
365                service_id: self.bind_service_reply.service_id,
366                method_id: Some(method_id),
367                args_proto: Some(encoded_args),
368                drop_reply: None,
369            }),
370        )
371    }
372
373    pub fn free_buffers(
374        &mut self,
375        current_task: &CurrentTask,
376        req: FreeBuffersRequest,
377    ) -> Result<u64, anyhow::Error> {
378        let method_id = self.method_id("FreeBuffers")?;
379        let mut encoded_args: Vec<u8> = Vec::with_capacity(req.encoded_len());
380        req.encode(&mut encoded_args)?;
381
382        self.send_message(
383            current_task,
384            ipc_frame::Msg::MsgInvokeMethod(ipc_frame::InvokeMethod {
385                service_id: self.bind_service_reply.service_id,
386                method_id: Some(method_id),
387                args_proto: Some(encoded_args),
388                drop_reply: None,
389            }),
390        )
391    }
392
393    pub fn next_frame_blocking(
394        &mut self,
395        current_task: &CurrentTask,
396    ) -> Result<IpcFrame, IpcReadError> {
397        self.frame_reader.next_frame_blocking(current_task)
398    }
399}
400
401// A perfetto compatible producer that can communicate with perfetto over a socket.
402// This struct provides a Rust compatible interface over the proto defined in
403// //third_party/perfetto/protos/perfetto/ipc/producer_port.proto.
404//
405// See https://perfetto.dev/docs/design-docs/api-and-abi#socket-protocol for details about the
406// protocol we implement here.
407pub struct Producer {
408    /// State for combining read byte data into IPC frames.
409    frame_reader: FrameReader,
410
411    /// Writer to write perfetto ipc to the socket
412    ipc_connection: IpcConnection,
413
414    /// After we connect, Perfetto will provide us with a service id which we will need to include
415    /// in all our messages going forwards so that Perfetto can identify us.
416    service_id: u32,
417
418    // When we connect to the perfetto socket, it informs us which method names correspond to which
419    // method ids. We save them here for reference when we need to invoke a method.
420    method_map: std::collections::HashMap<String, u32>,
421}
422
423impl Producer {
424    /// Opens a socket connection to the specified socket path and initializes the requisite
425    /// bookkeeping information.
426    pub fn new(current_task: &CurrentTask, socket: FileHandle) -> Result<Self, ProducerError> {
427        let mut producer = Self {
428            frame_reader: FrameReader::new(socket.clone()),
429            ipc_connection: IpcConnection::new(socket),
430            service_id: 0,
431            method_map: std::collections::HashMap::new(),
432        };
433
434        // The first thing we need to send over our newly connected socket is send a request to
435        // bind to the service.
436        producer.ipc_connection.bind_service("ProducerPort", current_task)?;
437
438        // Perfetto will then respond and tell us:
439        // 1) Our service_id, which we'll need to include in future messages to identify ourself
440        // 2) The list of methods that we can call using the "InvokeMethod" ipc
441        let reply_frame = producer.frame_reader.next_frame_blocking(current_task)?;
442
443        let ipc_frame::BindServiceReply { success, service_id, methods } = match reply_frame.msg {
444            Some(ipc_frame::Msg::MsgBindServiceReply(reply)) => reply,
445            m => {
446                return Err(ProducerError::InvalidResponse(format!(
447                    "Got unexpected reply message: {:?}",
448                    m
449                )));
450            }
451        };
452
453        if !success.unwrap_or(false) {
454            return Err(ProducerError::InvalidResponse("Bind to socket failed".into()));
455        }
456
457        // Build the possible methods we can call and save them for when we call InvokeMethod.
458        producer.method_map = methods
459            .into_iter()
460            .flat_map(|ipc_frame::bind_service_reply::MethodInfo { id, name }| match (id, name) {
461                (Some(id), Some(name)) => Some((name, id)),
462                _ => None,
463            })
464            .collect();
465        if let Some(service_id) = service_id {
466            producer.service_id = service_id
467        } else {
468            return Err(ProducerError::InvalidResponse(
469                "BindServiceReply did not include service_id".into(),
470            ));
471        }
472
473        Ok(producer)
474    }
475
476    /// Called once only after establishing the connection with the Service.
477    /// The service replies sending the shared memory file descriptor in reply.
478    pub fn initialize_connection(
479        &mut self,
480        request: InitializeConnectionRequest,
481        current_task: &CurrentTask,
482    ) -> Result<InitializeConnectionResponse, ProducerError> {
483        let (Some(reply), has_more) = self.invoke_method(
484            "InitializeConnection",
485            Some(request.encode_to_vec()),
486            current_task,
487        )?
488        else {
489            return Err(ProducerError::InvalidResponse("expected a response but got none".into()));
490        };
491        if has_more {
492            return Err(ProducerError::InvalidResponse(
493                "InitializeConnection should not stream but got a streaming response".into(),
494            ));
495        }
496        Ok(InitializeConnectionResponse::decode(reply.as_ref())?)
497    }
498
499    /// Advertises a new data source.
500    pub fn register_data_source(
501        &mut self,
502        request: RegisterDataSourceRequest,
503        current_task: &CurrentTask,
504    ) -> Result<RegisterDataSourceResponse, ProducerError> {
505        let (Some(reply), has_more) =
506            self.invoke_method("RegisterDataSource", Some(request.encode_to_vec()), current_task)?
507        else {
508            return Err(ProducerError::InvalidResponse(
509                "RegisterDataSource expected a response but got none".into(),
510            ));
511        };
512        if has_more {
513            return Err(ProducerError::InvalidResponse(
514                "RegisterDataSource should not stream but got a streaming response".into(),
515            ));
516        }
517        Ok(RegisterDataSourceResponse::decode(reply.as_ref())?)
518    }
519
520    /// Invoke a named method and block until the service replies.
521    fn invoke_method(
522        &mut self,
523        method_name: &str,
524        arguments: Option<Vec<u8>>,
525        current_task: &CurrentTask,
526    ) -> Result<(Option<Vec<u8>>, bool), InvokeMethodError> {
527        self.invoke_method_inner(method_name, arguments, current_task)?;
528
529        let reply_frame = self.frame_reader.next_frame_blocking(current_task)?;
530
531        let ipc_frame::InvokeMethodReply { success, has_more, reply_proto } = match reply_frame.msg
532        {
533            Some(ipc_frame::Msg::MsgInvokeMethodReply(reply)) => reply,
534            m => {
535                return Err(InvokeMethodError::InvalidResponse(format!(
536                    "unexpected reply message: {:?}",
537                    m
538                )));
539            }
540        };
541        match success {
542            Some(true) => Ok((reply_proto, has_more.unwrap_or(false))),
543            _ => {
544                return Err(InvokeMethodError::InvalidResponse(format!(
545                    "InvokeMethod Reply did not succeed. Reply: success: {:?}, has_more: {:?}, proto: {:?}",
546                    success, has_more, reply_proto,
547                )));
548            }
549        }
550    }
551
552    /// Invoke a named method without blocking for a response.
553    fn invoke_method_inner(
554        &mut self,
555        method_name: &str,
556        arguments: Option<Vec<u8>>,
557        current_task: &CurrentTask,
558    ) -> Result<(), InvokeMethodError> {
559        let Some(method_id) = self.method_map.get(method_name).copied() else {
560            return Err(InvokeMethodError::InvalidMethod(method_name.into()));
561        };
562        self.ipc_connection.invoke_method(self.service_id, method_id, arguments, current_task)?;
563        Ok(())
564    }
565
566    /// This is a backchannel to get asynchronous commands / notifications back
567    /// from the Service.
568    pub fn get_command_request(&mut self, current_task: &CurrentTask) -> Result<(), ProducerError> {
569        Ok(self.invoke_method_inner(
570            "GetAsyncCommand",
571            Some(GetAsyncCommandRequest {}.encode_to_vec()),
572            current_task,
573        )?)
574    }
575
576    /// After calling get_command_request, block until a response can be read.
577    pub fn get_command_response(
578        &mut self,
579        current_task: &CurrentTask,
580    ) -> Result<(Option<GetAsyncCommandResponse>, bool), ProducerError> {
581        let reply_frame = self.frame_reader.next_frame_blocking(current_task)?;
582
583        let ipc_frame::InvokeMethodReply { success, has_more, reply_proto } = match reply_frame.msg
584        {
585            Some(ipc_frame::Msg::MsgInvokeMethodReply(reply)) => reply,
586            m => {
587                return Err(ProducerError::InvalidResponse(format!(
588                    "Got unexpected reply message: {:?}",
589                    m
590                )));
591            }
592        };
593        if !success.unwrap_or(false) {
594            return Err(ProducerError::InvalidResponse(format!(
595                "InvokeMethod Reply did not include success. Reply: success: {:?}, has_more: {:?}, proto: {:?}",
596                success, has_more, reply_proto
597            )));
598        }
599
600        let Some(reply_proto) = reply_proto else {
601            return Err(ProducerError::InvalidResponse(
602                "InvokeMethod reply didn't include a proto".into(),
603            ));
604        };
605        Ok((
606            Some(GetAsyncCommandResponse::decode(reply_proto.as_ref())?),
607            has_more.unwrap_or(false),
608        ))
609    }
610}