Skip to main content

stress_mixed/
stress_mixed.rs

1use std::io;
2use std::time::{Duration, Instant};
3
4use cartel_sqlite::{Connection, params};
5
6const DEFAULT_OPERATION_COUNT: usize = 100_000;
7const DEFAULT_KEY_COUNT: usize = 1_024;
8const OPERATIONS_PER_CYCLE: usize = 5;
9const READ_OPERATIONS_PER_CYCLE: usize = 4;
10const WRITE_INCREMENT: i64 = 1;
11
12type ExampleError = Box<dyn std::error::Error>;
13
14struct WorkloadConfig {
15    operation_count: usize,
16    key_count: usize,
17}
18
19impl WorkloadConfig {
20    fn from_arguments() -> Result<Self, ExampleError> {
21        let mut arguments = std::env::args().skip(1);
22        let operation_count =
23            positive_count(arguments.next(), DEFAULT_OPERATION_COUNT, "operation count")?;
24        let key_count = positive_count(arguments.next(), DEFAULT_KEY_COUNT, "key count")?;
25        if let Some(unexpected_argument) = arguments.next() {
26            return Err(io::Error::new(
27                io::ErrorKind::InvalidInput,
28                format!("unexpected argument: {unexpected_argument}"),
29            )
30            .into());
31        }
32        i64::try_from(key_count).map_err(|_| {
33            io::Error::new(
34                io::ErrorKind::InvalidInput,
35                "key count exceeds SQLite INTEGER",
36            )
37        })?;
38        Ok(Self {
39            operation_count,
40            key_count,
41        })
42    }
43}
44
45struct WorkloadStats {
46    read_operations: usize,
47    write_operations: usize,
48    checksum: i64,
49    elapsed: Duration,
50}
51
52impl WorkloadStats {
53    fn operation_count(&self) -> usize {
54        self.read_operations + self.write_operations
55    }
56
57    fn throughput(&self) -> f64 {
58        if self.elapsed.is_zero() {
59            f64::INFINITY
60        } else {
61            self.operation_count() as f64 / self.elapsed.as_secs_f64()
62        }
63    }
64}
65
66fn positive_count(
67    argument: Option<String>,
68    default: usize,
69    name: &str,
70) -> Result<usize, ExampleError> {
71    let count = match argument {
72        Some(value) => value.parse::<usize>().map_err(|error| {
73            io::Error::new(
74                io::ErrorKind::InvalidInput,
75                format!("invalid {name} {value:?}: {error}"),
76            )
77        })?,
78        None => default,
79    };
80    if count == 0 {
81        return Err(io::Error::new(
82            io::ErrorKind::InvalidInput,
83            format!("{name} must be positive"),
84        )
85        .into());
86    }
87    Ok(count)
88}
89
90fn create_schema(connection: &Connection) -> cartel_sqlite::Result<()> {
91    connection.execute_batch(
92        "CREATE TABLE kv (
93            id INTEGER PRIMARY KEY,
94            value INTEGER NOT NULL
95        )",
96    )
97}
98
99fn seed_keys(connection: &mut Connection, key_count: usize) -> Result<(), ExampleError> {
100    let transaction = connection.transaction()?;
101    {
102        let mut insert_statement =
103            transaction.prepare("INSERT INTO kv (id, value) VALUES (?1, ?2)")?;
104        for key_index in 0..key_count {
105            let key = i64::try_from(key_index + 1)?;
106            insert_statement.execute(params![key, key])?;
107        }
108    }
109    transaction.commit()?;
110    Ok(())
111}
112
113fn execute_workload(
114    connection: &mut Connection,
115    config: &WorkloadConfig,
116) -> Result<WorkloadStats, ExampleError> {
117    let started = Instant::now();
118    let transaction = connection.transaction()?;
119    let mut read_operations = 0;
120    let mut write_operations = 0;
121    let mut checksum = 0_i64;
122    {
123        let mut point_lookup = transaction.prepare("SELECT value FROM kv WHERE id = ?1")?;
124        let mut point_update =
125            transaction.prepare("UPDATE kv SET value = value + ?2 WHERE id = ?1")?;
126        for operation_index in 0..config.operation_count {
127            let key_index = operation_index % config.key_count;
128            let key = i64::try_from(key_index + 1)?;
129            let cycle_index = operation_index % OPERATIONS_PER_CYCLE;
130            if cycle_index < READ_OPERATIONS_PER_CYCLE {
131                let value = point_lookup.query_row(params![key], |row| row.get::<_, i64>(0))?;
132                checksum = checksum.wrapping_add(value);
133                read_operations += 1;
134            } else {
135                let changed_rows = point_update.execute(params![key, WRITE_INCREMENT])?;
136                if changed_rows != 1 {
137                    return Err(io::Error::other(format!(
138                        "point update changed {changed_rows} rows for key {key}"
139                    ))
140                    .into());
141                }
142                write_operations += 1;
143            }
144        }
145    }
146    transaction.commit()?;
147    let elapsed = started.elapsed();
148    let stats = WorkloadStats {
149        read_operations,
150        write_operations,
151        checksum,
152        elapsed,
153    };
154    if stats.operation_count() != config.operation_count {
155        return Err(io::Error::other(format!(
156            "completed {} of {} workload operations",
157            stats.operation_count(),
158            config.operation_count
159        ))
160        .into());
161    }
162    Ok(stats)
163}
164
165fn main() -> Result<(), ExampleError> {
166    let config = WorkloadConfig::from_arguments()?;
167    let mut connection = Connection::open_in_memory()?;
168    create_schema(&connection)?;
169    seed_keys(&mut connection, config.key_count)?;
170    let stats = execute_workload(&mut connection, &config)?;
171
172    println!(
173        "sqlite mixed workload: operations={} reads={} writes={} keys={} elapsed_ms={:.3} throughput_ops_s={:.0} checksum={}",
174        stats.operation_count(),
175        stats.read_operations,
176        stats.write_operations,
177        config.key_count,
178        stats.elapsed.as_secs_f64() * 1_000.0,
179        stats.throughput(),
180        stats.checksum,
181    );
182    Ok(())
183}