Skip to main content

iggy_common/commands/system/
get_snapshot.rs

1/* Licensed to the Apache Software Foundation (ASF) under one
2 * or more contributor license agreements.  See the NOTICE file
3 * distributed with this work for additional information
4 * regarding copyright ownership.  The ASF licenses this file
5 * to you under the Apache License, Version 2.0 (the
6 * "License"); you may not use this file except in compliance
7 * with the License.  You may obtain a copy of the License at
8 *
9 *   http://www.apache.org/licenses/LICENSE-2.0
10 *
11 * Unless required by applicable law or agreed to in writing,
12 * software distributed under the License is distributed on an
13 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14 * KIND, either express or implied.  See the License for the
15 * specific language governing permissions and limitations
16 * under the License.
17 */
18
19use crate::BytesSerializable;
20use crate::Validatable;
21use crate::error::IggyError;
22use crate::{Command, GET_SNAPSHOT_FILE_CODE};
23use crate::{SnapshotCompression, SystemSnapshotType};
24use bytes::{BufMut, Bytes, BytesMut};
25use serde::{Deserialize, Serialize};
26use std::fmt::Display;
27use tracing::error;
28
29/// `GetSnapshot` command is used to get snapshot information.
30#[derive(Debug, Serialize, Deserialize, PartialEq)]
31pub struct GetSnapshot {
32    pub snapshot_types: Vec<SystemSnapshotType>,
33    pub compression: SnapshotCompression,
34}
35
36impl Default for GetSnapshot {
37    fn default() -> Self {
38        let types = vec![
39            SystemSnapshotType::FilesystemOverview,
40            SystemSnapshotType::ProcessList,
41            SystemSnapshotType::ResourceUsage,
42            SystemSnapshotType::ServerLogs,
43        ];
44        Self {
45            compression: SnapshotCompression::Deflated,
46            snapshot_types: types,
47        }
48    }
49}
50
51impl Command for GetSnapshot {
52    fn code(&self) -> u32 {
53        GET_SNAPSHOT_FILE_CODE
54    }
55}
56
57impl Validatable<IggyError> for GetSnapshot {
58    fn validate(&self) -> Result<(), IggyError> {
59        if self.snapshot_types.contains(&SystemSnapshotType::All) && self.snapshot_types.len() > 1 {
60            error!("When using 'All' snapshot type, no other types can be specified");
61            return Err(IggyError::InvalidCommand);
62        }
63        Ok(())
64    }
65}
66
67impl BytesSerializable for GetSnapshot {
68    fn to_bytes(&self) -> Bytes {
69        let mut bytes = BytesMut::new();
70        bytes.put_u8(self.compression.as_code());
71
72        bytes.put_u8(self.snapshot_types.len() as u8);
73        for snapshot_type in &self.snapshot_types {
74            bytes.put_u8(snapshot_type.as_code());
75        }
76
77        bytes.freeze()
78    }
79
80    fn from_bytes(bytes: Bytes) -> Result<GetSnapshot, IggyError> {
81        let mut index = 0;
82
83        let compression =
84            SnapshotCompression::from_code(*bytes.get(index).ok_or(IggyError::InvalidCommand)?)?;
85        index += 1;
86
87        let types_count = *bytes.get(index).ok_or(IggyError::InvalidCommand)? as usize;
88        index += 1;
89
90        let mut snapshot_types = Vec::with_capacity(types_count);
91
92        for _ in 0..types_count {
93            let tool =
94                SystemSnapshotType::from_code(*bytes.get(index).ok_or(IggyError::InvalidCommand)?)?;
95            index += 1;
96
97            snapshot_types.push(tool);
98        }
99
100        Ok(GetSnapshot {
101            compression,
102            snapshot_types,
103        })
104    }
105}
106
107impl Display for GetSnapshot {
108    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
109        write!(f, "GetSnapshotFile {{\n  snapshot_types: [\n")?;
110        for snapshot_type in &self.snapshot_types {
111            writeln!(f, "    {snapshot_type}")?;
112        }
113        write!(f, "  ],\n  Compression: {}\n}}", self.compression)
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use super::*;
120
121    #[test]
122    fn should_be_serialized_as_bytes() {
123        let get_snapshot_file_command = GetSnapshot {
124            compression: SnapshotCompression::Deflated,
125            snapshot_types: vec![SystemSnapshotType::FilesystemOverview],
126        };
127
128        let bytes = get_snapshot_file_command.to_bytes();
129
130        let deserialized = GetSnapshot::from_bytes(bytes.clone()).unwrap();
131
132        assert!(!bytes.is_empty());
133
134        assert_eq!(
135            deserialized.compression,
136            get_snapshot_file_command.compression
137        );
138        assert_eq!(
139            deserialized.snapshot_types,
140            get_snapshot_file_command.snapshot_types
141        );
142    }
143
144    #[test]
145    fn should_be_deserialized_from_bytes() {
146        let types = vec![SystemSnapshotType::FilesystemOverview];
147
148        let mut bytes = BytesMut::new();
149        bytes.put_u8(SnapshotCompression::Deflated.as_code());
150        bytes.put_u8(1);
151
152        bytes.put_u8(types.len() as u8);
153        for t in &types {
154            bytes.put_u8(t.as_code());
155        }
156
157        let get_snapshot = GetSnapshot::from_bytes(bytes.freeze());
158        assert!(get_snapshot.is_ok());
159
160        let get_snapshot = get_snapshot.unwrap();
161        assert_eq!(get_snapshot.snapshot_types.len(), 1);
162        assert_eq!(get_snapshot.compression, SnapshotCompression::Deflated);
163        assert_eq!(
164            get_snapshot.snapshot_types[0],
165            SystemSnapshotType::FilesystemOverview
166        );
167    }
168}