backend.rs 22.6 KB
Newer Older
1
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
Biswa Panda's avatar
Biswa Panda committed
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
// SPDX-License-Identifier: Apache-2.0

//! Backend
//!
//! An [`Backend`] is the final stage of the pipeline. It represents the execution of the LLM
//! on some processing hardware.
//!
//! At minimum, the Backend is split into two components, the [`Backend`] itself and a downstream [`ExecutionContext`].
//!
//! The [`ExecutionContext`] can be thought of as the core driver of the forward pass, whereas the [`Backend`] is the
//! manager of all resources and concurrent tasks surrounding the LLM execution context / forward pass.
//!
//! For almost every known scenario, detokenization and initial post processing must happen in the Backend.
//! Further post-processing can happen in the response stream. One example is the jailing mechanism for partial
//! hidden stop condition matches, which can be handled in the response stream rather than the backend.

18
use std::{collections::HashSet, sync::Arc, time::Instant};
Biswa Panda's avatar
Biswa Panda committed
19

20
use anyhow::Result;
Biswa Panda's avatar
Biswa Panda committed
21
22
use futures::stream::{self, StreamExt};

23
use crate::model_card::ModelDeploymentCard;
Neelay Shah's avatar
Neelay Shah committed
24
use dynamo_runtime::{
Biswa Panda's avatar
Biswa Panda committed
25
    pipeline::{
26
27
        AsyncEngineContextProvider, ManyOut, Operator, ResponseStream, ServerStreamingEngine,
        SingleIn, async_trait,
Biswa Panda's avatar
Biswa Panda committed
28
29
30
31
32
    },
    protocols::annotated::Annotated,
};

use crate::protocols::{
33
    TokenIdType,
Biswa Panda's avatar
Biswa Panda committed
34
    common::{
35
        StopConditions,
36
37
38
39
40
        llm_backend::{
            BackendOutput, EmbeddingsEngineOutput, FinishReason, LLMEngineOutput,
            PreprocessedRequest,
        },
        preprocessor::PreprocessedEmbeddingRequest,
41
        timing::RequestTracker,
Biswa Panda's avatar
Biswa Panda committed
42
43
    },
};
Nikita's avatar
Nikita committed
44
use crate::tokenizers::{DecodeStream, Tokenizer};
45
use dynamo_async_openai::types::StopReason;
Biswa Panda's avatar
Biswa Panda committed
46
47
48
49
50

/// Represents the output stream from the execution engine
pub type ExecutionOutputStream = Annotated<LLMEngineOutput>;

/// Context for executing LLM inference, engine consumes backend input and produces execution output stream
51
pub type ExecutionContext = ServerStreamingEngine<PreprocessedRequest, ExecutionOutputStream>;
Biswa Panda's avatar
Biswa Panda committed
52
53
54
55

/// Backend handles resource management and orchestrates LLM execution
#[allow(dead_code)]
pub struct Backend {
56
57
    pub tokenizer: Option<Tokenizer>, // Handles token encoding/decoding
    validate_engine_decode: bool,     // Enable validation of engine decoding
Biswa Panda's avatar
Biswa Panda committed
58
59
60
61
62
63
64
65
}

/// Internal state for managing token decoding and stream processing
#[allow(dead_code)]
struct DecoderUnfoldState {
    stream: ManyOut<ExecutionOutputStream>,
    decoder: Decoder,
    validate_engine_decode: bool,
66
67
    /// Set to true when a local stop condition is detected, causing the stream to end
    finished: bool,
Biswa Panda's avatar
Biswa Panda committed
68
69
70
}

impl Backend {
Nikita's avatar
Nikita committed
71
    pub fn from_tokenizer(tokenizer: Tokenizer) -> Arc<Self> {
72
        Arc::new(Self {
73
            tokenizer: Some(tokenizer),
Biswa Panda's avatar
Biswa Panda committed
74
            validate_engine_decode: false,
75
        })
Biswa Panda's avatar
Biswa Panda committed
76
77
    }

78
    pub fn from_mdc(mdc: &ModelDeploymentCard) -> Arc<Self> {
Nikita's avatar
Nikita committed
79
        match mdc.tokenizer() {
80
81
            Ok(tokenizer) => Self::from_tokenizer(tokenizer),
            Err(err) => {
Nikita's avatar
Nikita committed
82
                tracing::warn!(%err, "error loading tokenizer from ModelDeploymentCard");
83
                Arc::new(Self {
84
85
                    tokenizer: None,
                    validate_engine_decode: false,
86
                })
87
            }
88
        }
89
90
    }

Biswa Panda's avatar
Biswa Panda committed
91
92
93
    fn decoder(
        &self,
        stream: ManyOut<ExecutionOutputStream>,
94
        prompt_token_ids: &[TokenIdType],
Biswa Panda's avatar
Biswa Panda committed
95
        stop_conditions: StopConditions,
96
        skip_special_tokens: bool,
97
        include_stop_str_in_output: bool,
98
        tracker: Option<Arc<RequestTracker>>,
99
100
101
102
    ) -> anyhow::Result<DecoderUnfoldState> {
        let Some(tokenizer) = self.tokenizer.as_ref() else {
            anyhow::bail!("Backend built from blank ModelDeploymentCard, no tokenizer");
        };
103
        let decoder = Decoder::new(
104
            tokenizer.decode_stream(prompt_token_ids, skip_special_tokens),
105
            stop_conditions,
106
            include_stop_str_in_output,
107
            tracker,
108
        );
Biswa Panda's avatar
Biswa Panda committed
109

110
        Ok(DecoderUnfoldState {
Biswa Panda's avatar
Biswa Panda committed
111
112
113
            stream,
            decoder,
            validate_engine_decode: self.validate_engine_decode,
114
            finished: false,
115
        })
Biswa Panda's avatar
Biswa Panda committed
116
117
118
119
120
121
    }
}

#[async_trait]
impl
    Operator<
122
        SingleIn<PreprocessedRequest>,
Biswa Panda's avatar
Biswa Panda committed
123
        ManyOut<Annotated<BackendOutput>>,
124
        SingleIn<PreprocessedRequest>,
Biswa Panda's avatar
Biswa Panda committed
125
126
127
128
129
        ManyOut<Annotated<LLMEngineOutput>>,
    > for Backend
{
    async fn generate(
        &self,
130
131
        request: SingleIn<PreprocessedRequest>,
        next: ServerStreamingEngine<PreprocessedRequest, Annotated<LLMEngineOutput>>,
Biswa Panda's avatar
Biswa Panda committed
132
    ) -> Result<ManyOut<Annotated<BackendOutput>>> {
133
        let stop_conditions = request.stop_conditions.clone();
134
135
136

        let prompt_token_ids = request.token_ids.clone();

137
138
139
        // TODO: Consider updating default to true to match behavior of other frameworks
        let skip_special_tokens = request.output_options.skip_special_tokens.unwrap_or(false);

140
141
142
143
144
        // Extract include_stop_str_in_output from sampling_options (defaults to false)
        let include_stop_str_in_output = request
            .sampling_options
            .include_stop_str_in_output
            .unwrap_or(false);
145
        let tracker = request.tracker.clone();
146

Biswa Panda's avatar
Biswa Panda committed
147
148
149
        let next_stream = next.generate(request).await?;

        let context = next_stream.context();
150
151
152
153
154
        let state = self.decoder(
            next_stream,
            &prompt_token_ids,
            stop_conditions,
            skip_special_tokens,
155
            include_stop_str_in_output,
156
            tracker,
157
        )?;
Biswa Panda's avatar
Biswa Panda committed
158
159

        let processed_stream = stream::unfold(state, |mut state| async move {
160
161
162
163
164
            // If we've already detected a local stop condition, end the stream
            if state.finished {
                return None;
            }

Biswa Panda's avatar
Biswa Panda committed
165
166
167
168
169
170
171
172
173
174
175
            match state.stream.next().await {
                Some(output) => {
                    // move to state.process_output
                    // handle any error conditions / unwraps here

                    // events are pass thru
                    if output.is_event() || output.data.is_none() {
                        return Some((output, state));
                    }

                    // if we have a data field without an event, then we might need to update the data
176
177
178
179
180
                    if let Some(data) = &output.data
                        && data.text.is_some()
                        && !state.validate_engine_decode
                    {
                        return Some((output, state));
Biswa Panda's avatar
Biswa Panda committed
181
182
183
184
185
186
                    }

                    let data = output.data.as_ref().unwrap();

                    let result = state.decoder.process_token_ids(&data.token_ids).unwrap();

187
188
189
                    // NOTE: the `finish_reason` is computed from the generated `token_ids` alone.
                    // The `data` field can have a `finish_reason` set, coming from the underlying
                    // LLM inference `Engine`, and empty `token_ids`. See comment below for more details.
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
                    //
                    // stop_reason is only set for user-provided stop sequences, not for system
                    // EOS tokens (HiddenStopTokenDetected). This matches OpenAI API behavior where
                    // stop_reason is only present when a user-specified stop sequence is matched.
                    let (finish_reason, stop_reason) = match &result.stop_trigger {
                        Some(StopTrigger::MaxTokensLimit) => (Some(FinishReason::Length), None),
                        Some(StopTrigger::HiddenStopTokenDetected(_)) => {
                            // System EOS token - no stop_reason (user didn't request this stop)
                            (Some(FinishReason::Stop), None)
                        }
                        Some(StopTrigger::HiddenStopSequenceDetected(seq)) => {
                            // User-provided stop sequence (hidden from output)
                            (
                                Some(FinishReason::Stop),
                                Some(StopReason::String(seq.clone())),
                            )
Biswa Panda's avatar
Biswa Panda committed
206
                        }
207
208
209
210
211
212
213
214
                        Some(StopTrigger::VisibleStopSequenceDetected(seq)) => {
                            // User-provided stop sequence (included in output)
                            (
                                Some(FinishReason::Stop),
                                Some(StopReason::String(seq.clone())),
                            )
                        }
                        None => (None, None),
Biswa Panda's avatar
Biswa Panda committed
215
216
                    };

217
218
219
                    // If we detected a local stop condition, mark stream as finished
                    // so we stop iterating (upstream may keep generating, but we ignore it)
                    if finish_reason.is_some() && data.finish_reason.is_none() {
Biswa Panda's avatar
Biswa Panda committed
220
                        state.stream.context().stop_generating();
221
                        state.finished = true;
Biswa Panda's avatar
Biswa Panda committed
222
223
224
225
226
227
228
                    }

                    let text = result.text;
                    let tokens = result.tokens;

                    if state.validate_engine_decode {
                        if data.finish_reason != finish_reason {
229
                            tracing::warn!(
Biswa Panda's avatar
Biswa Panda committed
230
231
232
233
234
235
236
                                "finish reason mismatch: expected {:?}, got {:?}",
                                data.finish_reason,
                                finish_reason
                            );
                        }

                        if data.text.is_some() && data.text != text {
237
238
239
240
241
                            tracing::warn!(
                                "text mismatch: expected {:?}, got {:?}",
                                data.text,
                                text
                            );
Biswa Panda's avatar
Biswa Panda committed
242
243
244
245
246
247
248
                        }
                    }

                    // update output in-place
                    let mut output = output;
                    let mut data = output.data.take().unwrap();

249
250
251
252
253
254
255
256
                    // NOTE: If `finish_reason.is_some()`, then one of the stop conditions was triggered
                    // by the token generation. We should update the `data.finish_reason` in that case.
                    // However, if `finish_reason.is_none()`, it is possible that we are in the case where
                    // `data.token_ids` is empty, and `data.finish_reason` is already correctly set.
                    // In that case, `process_token_ids` above will rewrite `finish_reason` to `None`,
                    // which we don't want to propagate to `data.finish_reason`.
                    if finish_reason.is_some() {
                        data.finish_reason = finish_reason;
257
                        data.stop_reason = stop_reason;
258
                    }
Biswa Panda's avatar
Biswa Panda committed
259
260
261
262
263
264
265
266
267
268
                    data.text = text;
                    data.tokens = Some(tokens);

                    output.data = Some(data);

                    Some((output, state))
                }

                None => None,
            }
269
270
        })
        .fuse();
Biswa Panda's avatar
Biswa Panda committed
271
272

        // convert stream of processed Annotated<LLMEngineOutput> to Annotated<BackendOutput>
273
        //let mdcsum = self.mdcsum.clone();
Biswa Panda's avatar
Biswa Panda committed
274
275
276
277
278
279
280
281
        let stream = processed_stream.map(move |output| {
            output.map_data(|data| {
                Ok(BackendOutput {
                    token_ids: data.token_ids,
                    tokens: data.tokens.unwrap_or_default(),
                    text: data.text,
                    cum_log_probs: data.cum_log_probs,
                    log_probs: data.log_probs,
Greg Clark's avatar
Greg Clark committed
282
                    top_logprobs: data.top_logprobs,
Biswa Panda's avatar
Biswa Panda committed
283
                    finish_reason: data.finish_reason,
284
                    stop_reason: data.stop_reason,
285
                    //mdcsum: mdcsum.clone(),
286
                    index: data.index,
287
                    completion_usage: data.completion_usage,
288
                    disaggregated_params: data.disaggregated_params,
Biswa Panda's avatar
Biswa Panda committed
289
290
291
292
293
294
295
296
                })
            })
        });

        Ok(ResponseStream::new(Box::pin(stream), context))
    }
}

297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
#[async_trait]
impl
    Operator<
        SingleIn<PreprocessedEmbeddingRequest>,
        ManyOut<Annotated<EmbeddingsEngineOutput>>,
        SingleIn<PreprocessedEmbeddingRequest>,
        ManyOut<Annotated<EmbeddingsEngineOutput>>,
    > for Backend
{
    async fn generate(
        &self,
        request: SingleIn<PreprocessedEmbeddingRequest>,
        next: ServerStreamingEngine<
            PreprocessedEmbeddingRequest,
            Annotated<EmbeddingsEngineOutput>,
        >,
    ) -> Result<ManyOut<Annotated<EmbeddingsEngineOutput>>> {
        // For embeddings, we mostly pass through since no detokenization is needed
        // But we could add validation, logging, or other post-processing here
        let response_stream = next.generate(request).await?;

        // Could add embedding-specific post-processing here:
        // - Validation of embedding dimensions
        // - Normalization if requested
        // - Usage statistics validation

        Ok(response_stream)
    }
}

Biswa Panda's avatar
Biswa Panda committed
327
328
329
330
331
332
/// The [`Decoder`] object could be a member of either the internal LLM engine or part of the
/// postprocessor. If in the postprocessor, should be minimally in the same process or at very minimum
/// on the same physical machine connected by an IPC.
#[allow(dead_code)]
pub struct Decoder {
    decode_stream: DecodeStream,
333
    tracker: Option<Arc<RequestTracker>>,
Biswa Panda's avatar
Biswa Panda committed
334
335
336
337
338
339
340
341
342

    // do not trigger stop conditions until at least this many tokens have been generated
    min_tokens: u32,

    // single tokens that if found in the response will trigger a stop condition after the
    // minimum number of tokens have been generated
    hidden_stop_ids: HashSet<TokenIdType>,

    // text sequences that if found in the response will trigger a stop condition after the
343
    // minimum number of tokens have been generated (excluded from output)
Biswa Panda's avatar
Biswa Panda committed
344
345
    hidden_stop_sequences: Vec<String>,

346
347
348
349
    // text sequences that if found in the response will trigger a stop condition after the
    // minimum number of tokens have been generated (included in output)
    visible_stop_sequences: Vec<String>,

Biswa Panda's avatar
Biswa Panda committed
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
    // number of generated tokens
    generated_tokens: u32,

    // content jailed by partial hidden stop matches
    jail: String,

    // maximum number of bytes for the largest stop sequence
    jail_max_bytes: usize,

    // the number of bytes currently jailed
    jailed_bytes: usize,
}

#[allow(dead_code)]
#[derive(Debug)]
pub enum StopTrigger {
    MaxTokensLimit,
    HiddenStopTokenDetected(TokenIdType),
    HiddenStopSequenceDetected(String),
369
    VisibleStopSequenceDetected(String),
Biswa Panda's avatar
Biswa Panda committed
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
}

pub struct StepResult {
    pub token: Option<String>,
    pub stop_trigger: Option<StopTrigger>,
}

impl StepResult {
    fn ok(token: Option<String>) -> Self {
        Self {
            token,
            stop_trigger: None,
        }
    }

    fn with_stop_trigger(token: Option<String>, stop_trigger: StopTrigger) -> Self {
        Self {
            token,
            stop_trigger: Some(stop_trigger),
        }
    }
}

/// Result of processing a sequence of tokens
pub struct SeqResult {
    pub tokens: Vec<Option<String>>,       // Individual decoded tokens
    pub text: Option<String>,              // Combined decoded text
    pub stop_trigger: Option<StopTrigger>, // Reason for stopping generation, if any
}

#[allow(dead_code)]
impl Decoder {
    pub fn new(
        decode_stream: DecodeStream,
        stop_condition: StopConditions,
405
        include_stop_str_in_output: bool,
406
        tracker: Option<Arc<RequestTracker>>,
Biswa Panda's avatar
Biswa Panda committed
407
408
409
410
411
412
413
414
    ) -> Self {
        let hidden_stop_ids: HashSet<TokenIdType> = stop_condition
            .stop_token_ids_hidden
            .unwrap_or_default()
            .iter()
            .copied()
            .collect();

415
416
417
418
419
420
421
422
        // Categorize stop sequences based on include_stop_str_in_output:
        // - When true: user-provided stop sequences go to visible (included in output)
        // - When false: user-provided stop sequences go to hidden (excluded from output)
        let (hidden_stop_sequences, visible_stop_sequences) = if include_stop_str_in_output {
            (Vec::new(), stop_condition.stop.unwrap_or_default())
        } else {
            (stop_condition.stop.unwrap_or_default(), Vec::new())
        };
Biswa Panda's avatar
Biswa Panda committed
423

424
        // Calculate jail_max_bytes considering both hidden and visible stop sequences
Biswa Panda's avatar
Biswa Panda committed
425
426
        let jail_max_bytes = hidden_stop_sequences
            .iter()
427
            .chain(visible_stop_sequences.iter())
Biswa Panda's avatar
Biswa Panda committed
428
429
430
431
432
433
            .map(|x| x.len())
            .max()
            .unwrap_or(0);

        Self {
            decode_stream,
434
            tracker,
Biswa Panda's avatar
Biswa Panda committed
435
436
            hidden_stop_ids,
            hidden_stop_sequences,
437
            visible_stop_sequences,
Biswa Panda's avatar
Biswa Panda committed
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
            min_tokens: stop_condition.min_tokens.unwrap_or(0),
            generated_tokens: 0,
            jail: String::new(),
            jail_max_bytes,
            jailed_bytes: 0,
        }
    }

    /// Minimum amount of work to determine if a given generated/decoded sequence should be stopped
    /// This method can be called by the inner most loop of the LLM engine or minimally in the same
    /// process as the LLM engine.
    ///
    /// In the future, this method may kick off async cpu/tokio tasks and or async cuda tasks to
    /// handle logits post-processing and/or other tasks.
    pub fn step(&mut self, token_id: TokenIdType) -> Result<StepResult> {
        // increment the generated tokens
        self.generated_tokens += 1;

        // decode the token
457
        let detokenize_start = Instant::now();
Biswa Panda's avatar
Biswa Panda committed
458
        let token = self.decode_stream.step(token_id)?;
459
460
461
        if let Some(tracker) = &self.tracker {
            tracker.record_detokenize_latency(detokenize_start.elapsed());
        }
Biswa Panda's avatar
Biswa Panda committed
462
463
464
465
466
467
468
469
470

        // stop conditions to not apply until the minimum number of tokens have been generated
        if self.generated_tokens < self.min_tokens {
            return Ok(StepResult::ok(token));
        }

        // check for hidden stop tokens - eos takes precedence
        if self.hidden_stop_ids.contains(&token_id) {
            return Ok(StepResult::with_stop_trigger(
471
                None,
Biswa Panda's avatar
Biswa Panda committed
472
473
474
475
476
477
                StopTrigger::HiddenStopTokenDetected(token_id),
            ));
        }

        // check stop sequences - the jail will always hold at least the largest stop sequence
        // if jail_max_bytes is 0, then there are no stop sequences
478
479
480
481
482
483
        if self.jail_max_bytes > 0
            && let Some(token) = &token
        {
            let pre_append = self.jail.len();
            self.jail.push_str(token);

484
            // Check hidden stop sequences first (excluded from output)
485
486
487
            for seq in &self.hidden_stop_sequences {
                if let Some(offset) = galil_seiferas::gs_find(self.jail.as_bytes(), seq.as_bytes())
                {
488
                    // return only new bytes after pre_append .. offset (excluding stop sequence)
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
                    // example: seq = "ox", token = "boxes", return "b"
                    // note: this changes when we start jailing tokens for partial matches
                    // on the suffix of the jail with prefixes of the stop sequences
                    //
                    // we might have returned a partial match, if so, then offset < pre_append
                    // in that case, we return the empty string
                    let partial_token = if offset >= pre_append {
                        self.jail[pre_append..offset].to_string()
                    } else {
                        "".to_string()
                    };
                    return Ok(StepResult::with_stop_trigger(
                        Some(partial_token),
                        StopTrigger::HiddenStopSequenceDetected(seq.to_string()),
                    ));
Biswa Panda's avatar
Biswa Panda committed
504
                }
505
            }
Biswa Panda's avatar
Biswa Panda committed
506

507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
            // Check visible stop sequences (included in output)
            for seq in &self.visible_stop_sequences {
                if let Some(offset) = galil_seiferas::gs_find(self.jail.as_bytes(), seq.as_bytes())
                {
                    // For visible stop sequences, include the stop string in the output
                    // Return all text from pre_append up to and including the stop sequence
                    let stop_end = offset + seq.len();
                    let token_with_stop = if stop_end > pre_append {
                        self.jail[pre_append..stop_end].to_string()
                    } else {
                        // Stop sequence was entirely in previously returned text
                        "".to_string()
                    };
                    return Ok(StepResult::with_stop_trigger(
                        Some(token_with_stop),
                        StopTrigger::VisibleStopSequenceDetected(seq.to_string()),
                    ));
                }
            }

527
            Self::maybe_drain_to_max_bytes(&mut self.jail, self.jail_max_bytes);
Biswa Panda's avatar
Biswa Panda committed
528
529
530
531
532
533
534
        }

        Ok(StepResult::ok(token))
    }

    pub fn process_token_ids(&mut self, token_ids: &[TokenIdType]) -> Result<SeqResult> {
        let mut text: Option<String> = None;
535
        let mut tokens = Vec::with_capacity(token_ids.len());
Biswa Panda's avatar
Biswa Panda committed
536
537
538
539
540
541
542

        for token_id in token_ids {
            let StepResult {
                token,
                stop_trigger,
            } = self.step(*token_id)?;

543
544
            // Always include token text (for visible stops, the stop string is already in the token)
            if let Some(token) = &token {
545
546
                text.get_or_insert_with(|| String::with_capacity(token_ids.len()))
                    .push_str(token);
Biswa Panda's avatar
Biswa Panda committed
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
            }
            tokens.push(token);

            if let Some(stop_trigger) = stop_trigger {
                return Ok(SeqResult {
                    tokens,
                    text,
                    stop_trigger: Some(stop_trigger),
                });
            }
        }

        Ok(SeqResult {
            tokens,
            text,
            stop_trigger: None,
        })
    }

    fn jailed_string(&self) -> Option<String> {
        if self.jailed_bytes > 0 {
            // get the last jailed_bytes from the jail
            Some(self.jail[self.jail.len() - self.jailed_bytes..].to_string())
        } else {
            None
        }
    }
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597

    fn maybe_drain_to_max_bytes(s: &mut String, max_bytes: usize) {
        if s.len() > max_bytes {
            let mut drain_len = s.len() - max_bytes;
            while !s.is_char_boundary(drain_len) {
                drain_len -= 1;
            }
            s.drain(0..drain_len);
        }
    }
}

mod tests {

    #[test]
    fn test_char_boundary_drain() {
        use super::Decoder;
        let mut s = String::from("helloñworld"); // 12 bytes total ñ is 2 bytes
        let max_bytes = 6; // 12 - 6 = 6 which is inside ñ
        assert!(!s.is_char_boundary(s.len() - max_bytes)); // initially we are not on a char boundary
        Decoder::maybe_drain_to_max_bytes(&mut s, max_bytes);
        assert!(s.is_char_boundary(0)); // front of jail string on valid char boundary
        assert_eq!(s, "ñworld");
    }
Biswa Panda's avatar
Biswa Panda committed
598
}