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}