Skip to main content

reverie_process/
pid.rs

1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9use core::fmt;
10use core::hash::Hash;
11
12use serde::Deserialize;
13use serde::Serialize;
14
15/// A process ID (PID).
16#[derive(
17    Copy,
18    Clone,
19    Debug,
20    Eq,
21    PartialEq,
22    Ord,
23    PartialOrd,
24    Hash,
25    Serialize,
26    Deserialize
27)]
28pub struct Pid(libc::pid_t);
29
30impl Pid {
31    /// Creates `Pid` from a raw `pid_t`.
32    pub fn from_raw(pid: libc::pid_t) -> Self {
33        Self(pid)
34    }
35
36    /// Returns the PID of the calling process.
37    pub fn this() -> Self {
38        nix::unistd::Pid::this().into()
39    }
40
41    /// Returns the PID of the calling process.
42    pub fn parent() -> Self {
43        nix::unistd::Pid::parent().into()
44    }
45
46    /// Gets the raw `pid_t` from this `Pid`.
47    pub fn as_raw(self) -> libc::pid_t {
48        self.0
49    }
50
51    /// Returns a `Display`able that is color-coded. That is, the same PID will
52    /// get the same color. This makes it easy to visually recognize PIDs when
53    /// looking through logs.
54    ///
55    /// Note that while the same PIDs always have the same color, different PIDs
56    /// may also have the same color if they fall into the same color bucket.
57    pub fn colored(self) -> ColoredPid {
58        ColoredPid(self)
59    }
60}
61
62impl From<nix::unistd::Pid> for Pid {
63    fn from(pid: nix::unistd::Pid) -> Pid {
64        Self(pid.as_raw())
65    }
66}
67
68impl From<Pid> for nix::unistd::Pid {
69    fn from(pid: Pid) -> nix::unistd::Pid {
70        nix::unistd::Pid::from_raw(pid.as_raw())
71    }
72}
73
74impl From<Pid> for libc::pid_t {
75    fn from(pid: Pid) -> libc::pid_t {
76        pid.as_raw()
77    }
78}
79
80impl From<libc::pid_t> for Pid {
81    fn from(pid: libc::pid_t) -> Pid {
82        Pid::from_raw(pid)
83    }
84}
85
86impl fmt::Display for Pid {
87    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
88        fmt::Display::fmt(&self.0, f)
89    }
90}
91
92/// A colored pid.
93pub struct ColoredPid(Pid);
94
95impl ColoredPid {
96    /// Gets the ansi color code for the current PID. Returns `None` if not
97    /// writing to a terminal.
98    fn ansi_code(&self) -> Option<&'static str> {
99        if colored::control::SHOULD_COLORIZE.should_colorize() {
100            // Why not just use `colored::Colorize` you ask? It allocates a
101            // string in order to create the color code. Since we may log a lot
102            // of output that may contain a lot of PIDs, we don't want that to
103            // slow us down.
104            Some(match self.0.as_raw() % 14 {
105                0 => "\x1b[0;31m",  // Red
106                1 => "\x1b[0;32m",  // Green
107                2 => "\x1b[0;33m",  // Yellow
108                3 => "\x1b[0;34m",  // Blue
109                4 => "\x1b[0;35m",  // Magenta
110                5 => "\x1b[0;36m",  // Cyan
111                6 => "\x1b[0;37m",  // White
112                7 => "\x1b[1;31m",  // Bright red
113                8 => "\x1b[1;32m",  // Bright green
114                9 => "\x1b[01;33m", // Bright yellow
115                10 => "\x1b[1;34m", // Bright blue
116                11 => "\x1b[1;35m", // Bright magenta
117                12 => "\x1b[1;36m", // Bright cyan
118                _ => "\x1b[1;37m",  // Bright white
119            })
120        } else {
121            None
122        }
123    }
124}
125
126impl fmt::Display for ColoredPid {
127    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
128        if let Some(color) = self.ansi_code() {
129            write!(f, "{}{}\x1b[0m", color, self.0)
130        } else {
131            fmt::Display::fmt(&self.0, f)
132        }
133    }
134}