storage_stress_test_utils/
data.rs

1// Copyright 2020 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 rand::rngs::SmallRng;
6use rand::seq::SliceRandom;
7use rand::Rng;
8use std::cmp::min;
9
10/// Controls the compressibility of data generated for a file.
11/// This parameter is only meaningful for filesystems that support transparent compression.
12pub enum Compressibility {
13    Compressible,
14    Uncompressible,
15}
16
17// Controls the uncompressed size of the file
18pub enum UncompressedSize {
19    Exact(u64),
20    InRange(u64, u64),
21}
22
23/// Run lengths and bytes are sampled from a uniform distribution.
24/// Data created by this function achieves roughly 50% size reduction after compression.
25fn generate_compressible_data_bytes(rng: &mut SmallRng, size_bytes: u64) -> Vec<u8> {
26    let mut bytes: Vec<u8> = Vec::with_capacity(size_bytes as usize);
27
28    // The file is filled with compressible runs of one of the following bytes.
29    let byte_choices: [u8; 6] = [0xde, 0xad, 0xbe, 0xef, 0x42, 0x0];
30    let mut ptr = 0;
31    while ptr < size_bytes {
32        // A run is 10..1024 bytes long
33        let mut run_length = rng.gen_range(10..1024);
34
35        // In case the run goes past the file size
36        run_length = min(run_length, size_bytes - ptr);
37
38        // Decide whether this run should be compressible or not.
39        // This results in files that compress reasonably well (target 50% size reduction).
40        if rng.gen_bool(0.5) {
41            // Choose a byte for this run.
42            let choice = byte_choices.choose(rng).unwrap();
43
44            // Generate a run of compressible data.
45            for _ in 0..run_length {
46                bytes.push(*choice);
47            }
48        } else {
49            // Generate a run of random data.
50            for _ in 0..run_length {
51                bytes.push(rng.gen());
52            }
53        }
54
55        ptr += run_length;
56    }
57
58    // The file must be of the expected size
59    assert!(bytes.len() == size_bytes as usize);
60
61    bytes
62}
63
64/// Bytes are sampled from a uniform distribution.
65/// Data created by this function compresses badly.
66fn generate_uncompressible_data_bytes(rng: &mut SmallRng, size_bytes: u64) -> Vec<u8> {
67    let mut bytes: Vec<u8> = Vec::with_capacity(size_bytes as usize);
68    for _ in 0..size_bytes {
69        bytes.push(rng.gen());
70    }
71    bytes
72}
73
74/// A compact in-memory representation of data that can be stored in a file.
75pub struct FileFactory {
76    pub rng: SmallRng,
77    pub uncompressed_size: UncompressedSize,
78    pub compressibility: Compressibility,
79    pub file_name_count: u64,
80}
81
82impl FileFactory {
83    #[must_use]
84    pub fn new(
85        rng: SmallRng,
86        uncompressed_size: UncompressedSize,
87        compressibility: Compressibility,
88    ) -> Self {
89        Self { rng, uncompressed_size, compressibility, file_name_count: 0 }
90    }
91
92    /// Generate a unique filename that can be used for a file created by this factory.
93    pub fn generate_filename(&mut self) -> String {
94        self.file_name_count += 1;
95        format!("file_{}", self.file_name_count)
96    }
97
98    /// Generate all the bytes for this file in memory.
99    /// For a given FileData, this function is deterministic.
100    pub fn generate_bytes(&mut self) -> Vec<u8> {
101        let size_bytes = match self.uncompressed_size {
102            UncompressedSize::Exact(size_bytes) => size_bytes,
103            UncompressedSize::InRange(min, max) => self.rng.gen_range(min..max),
104        };
105
106        match self.compressibility {
107            Compressibility::Compressible => {
108                generate_compressible_data_bytes(&mut self.rng, size_bytes)
109            }
110            Compressibility::Uncompressible => {
111                generate_uncompressible_data_bytes(&mut self.rng, size_bytes)
112            }
113        }
114    }
115}
storage_stress_test_utils/data.rs

storage_stress_test_utils/
data.rs
