Refactoring resp_stream_*() (#704) (#856)

This originally started as refactoring designed to fix the quadratic memory growth observed in #704, but it's grown into a ground up rewrite to make it all easier to understand and considerably more efficient. Now all streaming functions that break data up into chunks (i.e. `resp_stream_lines()`, `resp_stream_sse()`, and `resp_stream_aws()`) use the same infrastructure, which has been extensively refactored to make it as easy as possible to understand.

Includes a number of small fixes to improve fidelity of AWS stream parsing.

Fixes #704

Author: hadley

DESCRIPTION

@@ -32,6 +32,7 @@ Suggests:
     bench,
     clipr,
     covr,
+    digest,
     docopt,
     httpuv,
     jose,

NEWS.md

@@ -15,6 +15,10 @@ e.g., `application/problem+json` (@cgiachalis, #782).
 * `req_body_form()` and `req_url_query()` no longer error with "C stack usage is too close to the limit" when given very long string values (#805).
 * `req_body_form()` now creates a valid empty request body when no parameters
   are provided (@arcresu, #836).
+* `resp_stream_aws()` now parses `byte`, `short`, and `integer` headers as signed integers, matching the AWS event-stream specification (previously they were incorrectly read as unsigned).
+* `resp_stream_lines()` no longer treats a bare carriage return (CR) as a line ending; only LF and CRLF terminate lines, which is what every modern streaming source produces.
+* `resp_stream_lines()` no longer warns when the stream ends without a final line terminator (which is routine when streaming), and its `warn` argument is (softly) deprecated.
+* `resp_stream_lines()`, `resp_stream_sse()`, and `resp_stream_aws()` now decode whole chunks at a time and hold the results in a queue, instead of rescanning and recopying the buffer for every line or event. This makes memory use and run time scale linearly rather than quadratically with the response size, so large streams use dramatically less memory and run much faster (e.g. reading a 1 MB response of short lines is now around 200x faster and allocates around 180x less memory) (#704).
 * `req_throttle()` can now enforce multiple rate limits at once: supply a vector to `capacity` (and `fill_time_s`) to create one token bucket per limit, and each request must satisfy all of them (#555).
 * `req_auth_aws_v4()` now correctly signs URLs containing encoded slashes (`%2F`) in path segments, such as ARNs in AWS Bedrock API paths (@thisisnic, #842).
 * `req_error()` is now applied to responses retrieved from the cache, so a custom `is_error` callback is respected on cache hits (#806).

R/req-perform-connection.R

@@ -252,3 +252,17 @@ StreamingBody <- R6::R6Class(
     conn = NULL
   )
 )
+
+# isOpen doesn't work for two reasons:
+# 1. It errors if con has been closed, rather than returning FALSE
+# 2. If returns TRUE if con has been closed and a new connection opened
+#
+# So instead we retrieve the connection from its number and compare to the
+# original connection. This works because connections have an undocumented
+# external pointer.
+isValid <- function(con) {
+  tryCatch(
+    identical(getConnection(con), con),
+    error = function(cnd) FALSE
+  )
+}

R/resp-stream-aws.R

@@ -1,17 +1,15 @@
 #' @export
 #' @rdname resp_stream_raw
-#' @order 2
+#' @order 4
 resp_stream_aws <- function(resp, max_size = Inf) {
-  event_bytes <- resp_boundary_pushback(
-    resp = resp,
-    max_size = max_size,
-    boundary_func = find_aws_event_boundary,
-    include_trailer = FALSE
-  )
+  splitter <- init_streaming_response(resp, AwsSplitter)
+  check_number_whole(max_size, min = 1, allow_infinite = TRUE)
 
-  if (is.null(event_bytes)) {
+  blocks <- stream_pull(resp, 1, splitter, max_size)
+  if (length(blocks) == 0L) {
     return()
   }
+  event_bytes <- blocks[[1L]]
 
   event <- parse_aws_event(event_bytes)
   if (resp_stream_show_body(resp)) {
@@ -26,24 +24,50 @@ resp_stream_aws <- function(resp, max_size = Inf) {
   event
 }
 
-find_aws_event_boundary <- function(buffer) {
-  # No valid AWS event message is less than 16 bytes
-  if (length(buffer) < 16) {
-    return(NULL)
-  }
-
-  # Read first 4 bytes as a big endian number
-  event_size <- parse_int(buffer[1:4])
-  if (event_size > length(buffer)) {
-    return(NULL)
+AwsSplitter <- R6::R6Class(
+  "AwsSplitter",
+  inherit = StreamSplitter,
+  public = list(
+    name = "resp_stream_aws()",
+    find_boundaries = function(buffer) find_aws_event_boundaries(buffer)
+  )
+)
+
+# Find every complete AWS event in a buffer by walking the 4-byte big-endian
+# length prefix at the start of each event. Returns a vector of split points
+# (the position one past the end of each complete event).
+find_aws_event_boundaries <- function(buffer) {
+  n <- length(buffer)
+  splits <- double()
+  pos <- 1
+  repeat {
+    # No valid AWS event message is less than 16 bytes.
+    if (n - pos + 1L < 16L) {
+      break
+    }
+    # Read the first 4 bytes of the event as a big endian number.
+    event_size <- parse_int(buffer[pos:(pos + 3L)])
+    if (event_size > n - pos + 1L) {
+      break
+    }
+    pos <- pos + event_size
+    splits[[length(splits) + 1L]] <- pos
   }
-
-  event_size + 1
+  splits
 }
 
-# Implementation from https://github.com/lifion/lifion-aws-event-stream/blob/develop/lib/index.js
-# This is technically buggy because it takes the header_length as a lower bound
-# but this shouldn't cause problems in practive
+# Parse a single AWS event-stream message (content type
+# application/vnd.amazon.eventstream). The binary format is documented by AWS:
+# * https://smithy.io/2.0/aws/amazon-eventstream.html (canonical protocol spec)
+# * https://docs.aws.amazon.com/lexv2/latest/dg/event-stream-encoding.html
+# Reference implementation: https://github.com/awslabs/aws-eventstream-java
+#
+# Key details: all integers are big-endian; the prelude (total + header lengths)
+# and the whole message each end in a GZIP/zlib CRC32; header value types
+# byte/short/integer/long are signed; timestamp is an int64 of epoch millis.
+#
+# We treat header_length as a lower bound rather than an exact count; this is
+# lenient but harmless and matches some reference implementations.
 parse_aws_event <- function(bytes) {
   i <- 1
   read_bytes <- function(n) {
@@ -80,9 +104,9 @@ parse_aws_event <- function(bytes) {
       type_enum(type),
       "TRUE" = TRUE,
       "FALSE" = FALSE,
-      BYTE = parse_int(read_bytes(1)),
-      SHORT = parse_int(read_bytes(2)),
-      INTEGER = parse_int(read_bytes(4)),
+      BYTE = parse_int(read_bytes(1), signed = TRUE),
+      SHORT = parse_int(read_bytes(2), signed = TRUE),
+      INTEGER = parse_int(read_bytes(4), signed = TRUE),
       LONG = parse_int64(read_bytes(8)),
       BYTE_ARRAY = read_bytes(length),
       CHARACTER = rawToChar(read_bytes(length)),
@@ -108,8 +132,13 @@ parse_aws_event <- function(bytes) {
 
 # Helpers ----------------------------------------------------------------
 
-parse_int <- function(x) {
-  sum(as.integer(x) * 256^rev(seq_along(x) - 1))
+parse_int <- function(x, signed = FALSE) {
+  v <- sum(as.integer(x) * 256^rev(seq_along(x) - 1))
+  if (signed && v >= 2^(8 * length(x) - 1)) {
+    # Interpret as two's complement.
+    v <- v - 2^(8 * length(x))
+  }
+  v
 }
 
 parse_int64 <- function(x) {
@@ -119,7 +148,7 @@ parse_int64 <- function(x) {
 }
 
 type_enum <- function(value) {
-  if (value < 0 || value > 10) {
+  if (value < 0 || value > 9) {
     cli::cli_abort("Unsupported type {value}.", .internal = TRUE)
   }
 
@@ -138,13 +167,6 @@ type_enum <- function(value) {
   )
 }
 
-hex_to_raw <- function(x) {
-  x <- gsub("(\\s|\n)+", "", x)
-
-  pairs <- substring(x, seq(1, nchar(x), by = 2), seq(2, nchar(x), by = 2))
-  as.raw(strtoi(pairs, 16L))
-}
-
 raw_to_hex <- function(x) {
   paste(as.character(x), collapse = "")
 }

R/resp-stream-lines.R

@@ -0,0 +1,57 @@
+#' @export
+#' @rdname resp_stream_raw
+#' @param lines The maximum number of lines to return at once.
+#' @param warn `r lifecycle::badge("deprecated")` `resp_stream_lines()` no longer
+#'   warns when the connection ends without a final EOL, so this argument is
+#'   ignored.
+#' @order 2
+resp_stream_lines <- function(
+  resp,
+  lines = 1,
+  max_size = Inf,
+  warn = deprecated()
+) {
+  splitter <- init_streaming_response(resp, LineSplitter)
+  check_number_whole(lines, min = 0, allow_infinite = TRUE)
+  check_number_whole(max_size, min = 1, allow_infinite = TRUE)
+  if (lifecycle::is_present(warn) && !isFALSE(warn)) {
+    lifecycle::deprecate_warn("1.2.3", "resp_stream_lines(warn)")
+  }
+
+  if (lines == 0) {
+    return(character())
+  }
+
+  encoding <- env_cache(resp$cache, "stream_encoding", resp_encoding(resp))
+  blocks <- stream_pull(resp, lines, splitter, max_size)
+  lines_read <- stream_parse_lines(blocks, encoding)
+  if (resp_stream_show_body(resp)) {
+    log_stream(lines_read)
+  }
+  lines_read
+}
+
+# Splits a stream into lines terminated by LF (and hence CRLF)
+LineSplitter <- R6::R6Class(
+  "LineSplitter",
+  inherit = StreamSplitter,
+  public = list(
+    name = "resp_stream_lines()",
+    find_boundaries = function(buffer) {
+      grepRaw(as.raw(0x0A), buffer, fixed = TRUE, all = TRUE) + 1L
+    },
+    # At end of stream, a trailing line without a terminator is still a line.
+    finish = function(remainder) {
+      if (length(remainder) == 0L) list() else list(remainder)
+    }
+  )
+)
+
+# Decode raw line blocks (each a line plus its trailing LF or CRLF) into a
+# character vector in `encoding`, dropping the terminators.
+stream_parse_lines <- function(blocks, encoding) {
+  text <- vapply(blocks, rawToChar, character(1))
+  Encoding(text) <- "bytes"
+  text <- iconv(text, encoding, "UTF-8")
+  sub("\r?\n$", "", text)
+}