Avoid write(::IO, ::String) invalidations from Stream and IO

* Avoid write(::IO, ::String) invalidations from Stream and deadline IO

Extending Base.write for String/AbstractString on a new IO type invalidates
every abstractly-inferred write(::IO, ::String) call site in the ecosystem
(~1380 instances each at using time, measured with @snoop_invalidations).
Hook Base.unsafe_write instead: Base's generic string write funnels through
it, so behavior is unchanged while the method-table impact disappears.

* Remove invalidation-prone abstract edges from write and websocket paths

- http1.jl / http_transport.jl: Int-assert length/write results where the
  stream is deliberately untyped, so loop arithmetic and comparisons stay
  concretely inferred instead of carrying >(::Any, ::Int) edges.
- http_server_streams.jl: write concrete Vector{UInt8}/String in the server
  write path; the abstract write(io, ::AbstractVector{UInt8}) and
  write(::Stream, ::AbstractString) edges get invalidated whenever any
  package defines write methods for its own types (LaTeXStrings, ...).
- http_websocket_codec.jl: use a Ref instead of a closure-reassigned local;
  the Core.Box capture infers Any and the resulting >(::UInt64, ::Any)
  edges get invalidated by packages defining broad comparison methods
  (SIMD.jl), recompiling the websocket receive path on first use.

* Define write(::Stream, ::UInt8) so generic AbstractStrings don't throw

String and SubString{String} writes route through unsafe_write, but Base's
fallbacks for other AbstractStrings (and Char) write char-by-char and bottom
out in write(io, ::UInt8), which previously threw "Stream does not support
byte I/O". Adding the byte primitive causes no invalidations (measured with
@snoop_invalidations against a loaded Bonito/Makie stack), since it is the
method every IO subtype is expected to define.

Author: SimonDanisch

src/http1.jl

@@ -597,15 +597,19 @@ end
 function _write_exact_bytes_body!(stream, body::BytesBody, expected_len::Int64)
     expected_len < 0 && throw(ArgumentError("expected_len must be >= 0"))
     expected_len == 0 && return nothing
-    available = (length(body.data) - body.next_index) + 1
+    # `body` is usually abstract here and `stream` is deliberately untyped;
+    # the Int asserts keep the arithmetic/comparisons concretely inferred so
+    # this method carries no invalidation-prone `>=(::Any, ::Int)` edges
+    len = length(body.data)::Int
+    available = (len - body.next_index) + 1
     available >= expected_len || throw(ProtocolError("body ended before expected Content-Length bytes were written"))
     stop_index = body.next_index + Int(expected_len) - 1
-    chunk = if body.next_index == 1 && stop_index == length(body.data)
+    chunk = if body.next_index == 1 && stop_index == len
         body.data
     else
         view(body.data, body.next_index:stop_index)
     end
-    n = write(stream, chunk)
+    n = Int(write(stream, chunk))
     n == expected_len || throw(ProtocolError("transport short write"))
     body.next_index = stop_index + 1
     return nothing

src/http_server_streams.jl

@@ -295,7 +295,10 @@ function _server_write(stream::Stream, data::AbstractVector{UInt8})::Int
     if stream.write_mode == _ServerStreamWriteMode.CHUNKED
         io = IOBuffer()
         print(io, string(length(data), base=16), "\r\n")
-        write(io, data)
+        # concrete bytes: `write(io, ::AbstractVector{UInt8})` is an
+        # invalidation-prone edge (any package adding write methods recompiles
+        # the server write path); no copy in the common Vector case
+        write(io, data isa Vector{UInt8} ? data : Vector{UInt8}(data))
         write(io, "\r\n")
         _write_server_stream_bytes!(stream, take!(io))
     else
@@ -394,11 +397,15 @@ function _write_response_body_to_stream!(stream::Stream, body)::Nothing
         return nothing
     end
     if body isa AbstractString
-        write(stream, body::AbstractString)
+        # concrete String: a `write(::Stream, ::AbstractString)` edge gets
+        # invalidated whenever any package defines a write method for its
+        # own string type (LaTeXStrings, ...), recompiling the server
+        # response path
+        write(stream, body isa String ? body : String(body))
         return nothing
     end
     if body isa AbstractVector{UInt8}
-        write(stream, body::AbstractVector{UInt8})
+        write(stream, body isa Vector{UInt8} ? body : Vector{UInt8}(body))
         return nothing
     end
     if body isa AbstractBody

src/http_stream.jl

@@ -253,16 +253,34 @@ function write(stream::Stream{true}, data::AbstractVector{UInt8})::Int
 end
 write(stream::Stream{false}, data::AbstractVector{UInt8}) = _server_write(stream, data)
 
-function _client_stream_write(stream::Stream{true}, data::AbstractString)::Int
+# Strings are handled via unsafe_write instead of write(::Stream, ::AbstractString)
+# methods: Base's generic `write(io::IO, s::Union{String, SubString{String}})`
+# funnels through unsafe_write, and extending `write` for string types on a new
+# IO type invalidates every abstractly-inferred `write(::IO, ::String)` call
+# site in the ecosystem (measured: ~1380 invalidated instances at `using` time).
+function Base.unsafe_write(stream::Stream{true}, p::Ptr{UInt8}, n::UInt)::UInt
+    n == 0 && return UInt(0)
     (@atomic :acquire stream.started) && throw(ArgumentError("cannot write request body after response reading has started"))
     (@atomic :acquire stream.write_closed) && throw(ArgumentError("request body writes are closed"))
-    return write(stream.request_buffer, String(data))
+    return unsafe_write(stream.request_buffer, p, n)
+end
+
+function Base.unsafe_write(stream::Stream{false}, p::Ptr{UInt8}, n::UInt)::UInt
+    n == 0 && return UInt(0)
+    data = Vector{UInt8}(undef, Int(n))
+    GC.@preserve data unsafe_copyto!(pointer(data), p, n)
+    return UInt(_server_write(stream, data))
 end
 
-write(stream::Stream{true}, data::AbstractString)::Int = _client_stream_write(stream, data)
-write(stream::Stream{true}, data::Union{String,SubString{String}})::Int = _client_stream_write(stream, data)
-write(stream::Stream{false}, data::AbstractString) = _server_write(stream, data)
-write(stream::Stream{false}, data::Union{String,SubString{String}}) = _server_write(stream, data)
+# The byte-I/O primitive: Base's fallbacks for `write(io, ::AbstractString)`
+# and `write(io, ::Char)` bottom out in `write(io, ::UInt8)`, so without this
+# generic strings throw "does not support byte I/O".
+function write(stream::Stream{true}, b::UInt8)::Int
+    (@atomic :acquire stream.started) && throw(ArgumentError("cannot write request body after response reading has started"))
+    (@atomic :acquire stream.write_closed) && throw(ArgumentError("request body writes are closed"))
+    return write(stream.request_buffer, b)
+end
+write(stream::Stream{false}, b::UInt8) = _server_write(stream, UInt8[b])
 
 function closewrite(stream::Stream{true})
     @atomic :release stream.write_closed = true

src/http_transport.jl

@@ -78,9 +78,14 @@ Base.write(io::_RequestDeadlineWriteIO, data::Vector{UInt8})::Int = _request_dea
 Base.write(io::_RequestDeadlineWriteIO, data::StridedVector{UInt8})::Int = _request_deadline_write_bytes!(io, data)
 Base.write(io::_RequestDeadlineWriteIO, data::AbstractVector{UInt8})::Int = _request_deadline_write_bytes!(io, data)
 
-function Base.write(io::_RequestDeadlineWriteIO, data::Union{String,SubString{String}})::Int
+# Hook unsafe_write instead of write(::IO, ::Union{String, SubString{String}}):
+# Base's generic string write funnels through unsafe_write, so strings still
+# get the deadline applied, while extending `write` for String on a new IO
+# type invalidates every abstractly-inferred `write(::IO, ::String)` call site
+# in the ecosystem (measured: 1383 invalidated instances at `using` time).
+function Base.unsafe_write(io::_RequestDeadlineWriteIO, p::Ptr{UInt8}, n::UInt)::Int
     _apply_request_write_deadline!(io)
-    return write(io.inner, data)
+    return unsafe_write(io.inner, p, n)
 end
 
 @inline function _fill_conn_reader!(reader::_ConnReader)::Int
@@ -430,7 +435,10 @@ function _write_exact_bytes_body_transport!(
         else
             view(body.data, body.next_index:stop_index)
         end
-        n = write(stream, chunk)
+        # `stream` is deliberately untyped (multiple transports); without the
+        # Int assert `n` infers Any and poisons the loop comparisons with
+        # invalidation-prone `>(::Any, ::Int)` edges
+        n = Int(write(stream, chunk))
         n == chunk_len || throw(ProtocolError("transport short write"))
         body.next_index = stop_index + 1
         remaining -= n

test/http_server_http1_tests.jl

@@ -948,6 +948,44 @@ end
     end
 end
 
+@testset "HTTP stream byte I/O supports generic AbstractStrings and Chars" begin
+    # Strings are written via unsafe_write; generic AbstractStrings and Chars
+    # fall back to Base's per-byte path, which requires write(io, ::UInt8).
+    server = HT.listen!("127.0.0.1", 0; listenany = true) do stream
+        _ = HT.startread(stream)
+        body = String(read(stream))
+        HT.setstatus(stream, 200)
+        HT.startwrite(stream)
+        if isempty(body)
+            write(stream, Test.GenericString("generic"))
+            write(stream, ' ')
+            write(stream, 0x21)
+        else
+            write(stream, body)
+        end
+        return nothing
+    end
+    address = HT.server_addr(server)
+    try
+        resp = HT.get("http://$(address)/")
+        @test resp.status == 200
+        @test String(resp.body) == "generic !"
+
+        echoed = Ref("")
+        resp2 = HT.open(:POST, "http://$(address)/") do io
+            write(io, Test.GenericString("client generic"))
+            write(io, '!')
+            _ = HT.startread(io)
+            echoed[] = String(read(io))
+        end
+        @test resp2.status == 200
+        @test echoed[] == "client generic!"
+    finally
+        _run_with_timeout(() -> HT.forceclose(server); label = "server forceclose")
+        _run_with_timeout(() -> wait(server); label = "server task completion")
+    end
+end
+
 @testset "HTTP server stream handlers reject fixed-length mismatches before writing malformed bodies" begin
     server = HT.listen!("127.0.0.1", 0; listenany = true) do stream
             request = HT.startread(stream)