Skip to content

Commit

Permalink
Update readbytes! to avoid using internal API of IOBuffer less (#1148)
Browse files Browse the repository at this point in the history
* Update readbytes! to avoid using internal API of IOBuffer less

Note that this claims to be supported for all GenericIOBuffer, but Base only defines this function on a normal IOBuffer for which `pointer(data)` is contiguous space

* Update Streams.jl

* Update Streams.jl

* Update Streams.jl

* fix some tests relying on implementation details
  • Loading branch information
vtjnash authored Mar 5, 2024
1 parent d21ae94 commit e27493d
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 17 deletions.
19 changes: 7 additions & 12 deletions src/Streams.jl
Original file line number Diff line number Diff line change
Expand Up @@ -279,21 +279,15 @@ function Base.unsafe_read(http::Stream, p::Ptr{UInt8}, n::UInt)
nothing
end

@noinline function bufcheck(buf::Base.GenericIOBuffer, n)
requested_buffer_capacity = (buf.append ? buf.size : (buf.ptr - 1)) + n
requested_buffer_capacity > length(buf.data) && throw(ArgumentError("Unable to grow response stream IOBuffer $(length(buf.data)) large enough for response body size: $requested_buffer_capacity"))
end

function Base.readbytes!(http::Stream, buf::Base.GenericIOBuffer, n=bytesavailable(http))
Base.ensureroom(buf, n)
# check if there's enough room in buf to write n bytes
bufcheck(buf, n)
data = buf.data
GC.@preserve data unsafe_read(http, pointer(data, (buf.append ? buf.size + 1 : buf.ptr)), n)
p, nbmax = Base.alloc_request(buf, UInt(n))
nbmax < n && throw(ArgumentError("Unable to grow response stream IOBuffer $nbmax large enough for response body size: $n"))
GC.@preserve buf unsafe_read(http, p, UInt(n))
# TODO: use `Base.notify_filled(buf, Int(n))` here, but only once it is identical to this:
if buf.append
buf.size += n
buf.size += Int(n)
else
buf.ptr += n
buf.ptr += Int(n)
buf.size = max(buf.size, buf.ptr - 1)
end
return n
Expand Down Expand Up @@ -327,6 +321,7 @@ function IOExtras.readuntil(http::Stream, f::Function)
update_ntoread(http, length(bytes))
return bytes
catch ex
ex isa EOFError || rethrow()
# if we error, it means we didn't find what we were looking for
# TODO: this seems very sketchy
return UInt8[]
Expand Down
13 changes: 8 additions & 5 deletions test/client.jl
Original file line number Diff line number Diff line change
Expand Up @@ -140,7 +140,7 @@ end
# wrapping pre-allocated buffer in IOBuffer will write to buffer directly
io = IOBuffer(body; write=true)
r = HTTP.get("https://$httpbin/bytes/100"; response_stream=io, socket_type_tls=tls)
@test body === r.body.data
@test Base.mightalias(body, r.body.data)

# if provided buffer is too small, we won't grow it for user
body = zeros(UInt8, 10)
Expand All @@ -156,27 +156,30 @@ end
# but if you wrap it in a writable IOBuffer, we will grow it
io = IOBuffer(body; write=true)
r = HTTP.get("https://$httpbin/bytes/100"; response_stream=io, socket_type_tls=tls)
# same Array, though it was resized larger
@test body === r.body.data
# might be a new Array, resized larger
body = take!(io)
@test length(body) == 100

# and you can reuse it
seekstart(io)
r = HTTP.get("https://$httpbin/bytes/100"; response_stream=io, socket_type_tls=tls)
# same Array, though it was resized larger
@test body === r.body.data
# `take!` should have given it a new Array
@test !Base.mightalias(body, r.body.data)
body = take!(io)
@test length(body) == 100

# we respect ptr and size
body = zeros(UInt8, 100)
io = IOBuffer(body; write=true, append=true) # size=100, ptr=1
r = HTTP.get("https://$httpbin/bytes/100"; response_stream=io, socket_type_tls=tls)
body = take!(io)
@test length(body) == 200

body = zeros(UInt8, 100)
io = IOBuffer(body, write=true, append=false)
write(io, body) # size=100, ptr=101
r = HTTP.get("https://$httpbin/bytes/100"; response_stream=io, socket_type_tls=tls)
body = take!(io)
@test length(body) == 200

end
Expand Down

0 comments on commit e27493d

Please sign in to comment.