sachy.dev
A pit full of rusty nails
1//! Crate for defining a HTML generator based on a markov chain source, using a string
2//! interner to reduce memory usage both within a markov chain and across multiple chains.
3//!
4
5use core::task::Poll;
6use std::{
7 path::Path,
8 pin::Pin,
9 sync::Arc,
10 time::{Duration, Instant},
11};
12
13use axum::extract::MatchedPath;
14use bytes::{Bytes, BytesMut};
15use color_eyre::Result;
16use futures_lite::Stream;
17use nailbox::{boxed_future_within, try_arc_within};
18use nailconfig::NailConfig;
19use nailkov::NailKov;
20use nailrng::FastRng;
21use pin_project_lite::pin_project;
22use tokio::time::Sleep;
23
24use crate::{
25 delay::delay_output,
26 html_gen::{
27 extra, footer, initial_content, main_content, static_content, static_title, text_generator,
28 },
29};
30
31pub use crate::template::*;
32
33mod delay;
34mod html_gen;
35mod template;
36
37#[derive(Clone)]
38pub struct MarkovGen {
39 chain: Arc<NailKov>,
40}
41
42pin_project! {
43 #[project = GeneratorStateProj]
44 enum GeneratorState {
45 Template,
46 Content,
47 GeneratingContent {
48 handle: Pin<Box<dyn Future<Output = Bytes> + Send>>,
49 keep_generating: bool,
50 },
51 Delay {
52 delay: Pin<Box<Sleep>>,
53 },
54 Finished,
55 }
56}
57
58pin_project! {
59 pub struct MarkovStream {
60 path: MatchedPath,
61 config: Arc<NailConfig>,
62 markov: MarkovGen,
63 start_time: Instant,
64 total_bytes: usize,
65 template: Template,
66 cursor: TemplateCursor,
67 page_title: Option<Box<[u8]>>,
68 rng: FastRng,
69 #[pin]
70 state: GeneratorState,
71 }
72}
73
74impl MarkovStream {
75 pub fn new(
76 markov: MarkovGen,
77 path: MatchedPath,
78 config: Arc<NailConfig>,
79 template: Template,
80 rng: FastRng,
81 ) -> Self {
82 Self {
83 path,
84 config,
85 markov,
86 total_bytes: 0,
87 start_time: Instant::now(),
88 state: GeneratorState::Template,
89 cursor: TemplateCursor::new(template.get_template()),
90 rng,
91 template,
92 page_title: None,
93 }
94 }
95}
96
97impl Stream for MarkovStream {
98 type Item = Bytes;
99
100 #[cfg_attr(
101 feature = "detailed_traces",
102 tracing::instrument(level = "trace", name = "MarkovStream::poll_next", skip_all)
103 )]
104 #[inline]
105 fn poll_next(
106 mut self: std::pin::Pin<&mut Self>,
107 cx: &mut std::task::Context<'_>,
108 ) -> Poll<Option<Self::Item>> {
109 let mut this = self.as_mut().project();
110
111 'outer: loop {
112 let mut buffer = BytesMut::new();
113
114 match this.state.as_mut().project() {
115 GeneratorStateProj::Template => 'inner: loop {
116 match this.cursor.write_template(&mut buffer) {
117 template::TemplateState::Title => {
118 let title = this.page_title.get_or_insert_with(|| {
119 this.template.get_static_content().map_or_else(
120 || {
121 text_generator(&this.markov.chain, 24, this.rng)
122 .copied()
123 .collect()
124 },
125 |title| static_title(title).copied().collect(),
126 )
127 });
128
129 *this.total_bytes += title.len();
130
131 buffer.extend_from_slice(title);
132
133 continue 'inner;
134 }
135 template::TemplateState::Initial => {
136 let handle = boxed_future_within(|| {
137 initial_content(
138 buffer,
139 this.markov.chain.clone(),
140 this.config.clone(),
141 this.rng.fork(),
142 )
143 });
144
145 this.state.set(GeneratorState::GeneratingContent {
146 handle,
147 keep_generating: false,
148 });
149
150 continue 'outer;
151 }
152 template::TemplateState::Main => {
153 if let Some(content) = this.template.get_static_content() {
154 let len = buffer.len();
155
156 buffer.extend(static_content(content));
157
158 this.state.set(GeneratorState::Template);
159
160 let len = buffer.len() - len;
161
162 *this.total_bytes += len;
163
164 continue 'inner;
165 } else {
166 this.state.set(GeneratorState::Content);
167
168 continue 'outer;
169 }
170 }
171 template::TemplateState::Extra => {
172 let bytes = extra(&mut buffer, this.config, this.rng);
173
174 *this.total_bytes += bytes;
175
176 continue 'inner;
177 }
178 template::TemplateState::Footer => {
179 let handle = boxed_future_within(|| {
180 footer(
181 buffer,
182 this.markov.chain.clone(),
183 this.path.clone(),
184 this.config.clone(),
185 this.rng.fork(),
186 )
187 });
188
189 this.state.set(GeneratorState::GeneratingContent {
190 handle,
191 keep_generating: false,
192 });
193
194 continue 'outer;
195 }
196 template::TemplateState::Finished => {
197 let elapsed_time = this.start_time.elapsed().as_micros();
198
199 tracing::trace!(
200 "payload.size" = *this.total_bytes,
201 "duration.us" = elapsed_time,
202 "Stream finished in {:.2}ms",
203 (elapsed_time as f32) * 1e-3
204 );
205
206 this.state.set(GeneratorState::Finished);
207
208 return Poll::Ready(Some(buffer.freeze()));
209 }
210 }
211 },
212 GeneratorStateProj::Content => {
213 let time_limit = Duration::from_secs(this.config.generator.timeout);
214
215 if time_limit.as_secs() > 0
216 && this.start_time.elapsed().as_secs() >= time_limit.as_secs()
217 {
218 this.state.set(GeneratorState::Template);
219 continue 'outer;
220 }
221
222 if *this.total_bytes >= (this.config.generator.payload_size * 1024) {
223 this.state.set(GeneratorState::Template);
224 continue 'outer;
225 }
226
227 let handle = boxed_future_within(|| {
228 main_content(
229 buffer,
230 this.markov.chain.clone(),
231 this.config.clone(),
232 this.rng.fork(),
233 )
234 });
235
236 this.state.set(GeneratorState::GeneratingContent {
237 handle,
238 keep_generating: true,
239 });
240 }
241 GeneratorStateProj::GeneratingContent {
242 handle,
243 keep_generating,
244 } => match handle.as_mut().poll(cx) {
245 Poll::Ready(content) => {
246 *this.total_bytes += content.len();
247
248 if *keep_generating {
249 if let Some(delay) = delay_output(this.config) {
250 this.state.set(GeneratorState::Delay { delay });
251 } else {
252 this.state.set(GeneratorState::Content);
253 }
254 } else {
255 this.state.set(GeneratorState::Template);
256 }
257
258 return Poll::Ready(Some(content));
259 }
260 Poll::Pending => return Poll::Pending,
261 },
262 GeneratorStateProj::Delay { delay } => match delay.as_mut().poll(cx) {
263 Poll::Ready(_) => {
264 this.state.set(GeneratorState::Content);
265
266 continue;
267 }
268 Poll::Pending => return Poll::Pending,
269 },
270 GeneratorStateProj::Finished => {
271 return Poll::Ready(None);
272 }
273 }
274 }
275 }
276}
277
278impl MarkovGen {
279 pub fn new(input: impl AsRef<Path>) -> Result<Self> {
280 let file = std::fs::read_to_string(input.as_ref())?;
281
282 let chain = try_arc_within(|| NailKov::from_input(&file))?;
283
284 Ok(Self { chain })
285 }
286
287 #[inline]
288 pub fn into_stream(
289 self,
290 path: MatchedPath,
291 config: Arc<NailConfig>,
292 template: Template,
293 rng: FastRng,
294 ) -> MarkovStream {
295 MarkovStream::new(self, path, config, template, rng)
296 }
297}