Add ability to log upload rates (#39)

Author: quinnj

src/blobs.jl

@@ -46,8 +46,8 @@ API.startMultipartUpload(x::Container, key; kw...) = nothing
 
 function API.uploadPart(x::Container, url, part, partNumber, uploadId; kw...)
     blockid = base64encode(lpad(partNumber - 1, 64, '0'))
-    Azure.put(url, [], part; query=Dict("comp" => "block", "blockid" => blockid), kw...)
-    return blockid
+    resp = Azure.put(url, [], part; query=Dict("comp" => "block", "blockid" => blockid), kw...)
+    return (blockid, Base.get(resp.request.context, :nbytes_written, 0))
 end
 
 function API.completeMultipartUpload(x::Container, url, eTags, uploadId; kw...)

src/get.jl

@@ -63,7 +63,8 @@ function getObjectImpl(x::AbstractStore, key::String, out::ResponseBodyType=noth
     allowMultipart::Bool=true,
     decompress::Bool=false,
     zlibng::Bool=false,
-    headers=HTTP.Headers(), kw...)
+    headers=HTTP.Headers(),
+    lograte::Bool=false, kw...)
 
     # keyword arg handling
     if (out isa AbstractVector{UInt8} && length(out) <= multipartThreshold)
@@ -170,6 +171,6 @@ function getObjectImpl(x::AbstractStore, key::String, out::ResponseBodyType=noth
     end_time = time()
     bytes = nbytes[]
     gbits_per_second = bytes == 0 ? 0 : (((8 * bytes) / 1e9) / (end_time - start_time))
-    @debug "CloudStore.get complete with bandwidth: $(gbits_per_second) Gbps"
+    lograte && @info "CloudStore.get complete with bandwidth: $(gbits_per_second) Gbps"
     return res
 end

src/put.jl

@@ -52,13 +52,18 @@ function putObjectImpl(x::AbstractStore, key::String, in::RequestBodyType;
     batchSize::Int=defaultBatchSize(),
     allowMultipart::Bool=true,
     zlibng::Bool=false,
-    compress::Bool=false, credentials=nothing, kw...)
+    compress::Bool=false, credentials=nothing,
+    lograte::Bool=false, kw...)
 
+    start_time = time()
     N = nbytes(in)
+    wbytes = Threads.Atomic{Int}(0)
     if N <= multipartThreshold || !allowMultipart
         body = prepBody(in, compress, zlibng)
         resp = putObject(x, key, body; credentials, kw...)
-        return Object(x, credentials, key, N, etag(HTTP.header(resp, "ETag")))
+        wbytes[] = get(resp.request.context, :nbytes_written, 0)
+        obj = Object(x, credentials, key, N, etag(HTTP.header(resp, "ETag")))
+        @goto done
     end
     # multipart upload
     uploadState = startMultipartUpload(x, key; credentials, kw...)
@@ -76,15 +81,13 @@ function putObjectImpl(x::AbstractStore, key::String, in::RequestBodyType;
             eof(body) && break
             n = (j - 1) * batchSize + i
             part = _read(body, partSize)
-            let n=n, part=part
-                Threads.@spawn begin
-                    eTag = uploadPart(x, url, part, n, uploadState; credentials, kw...)
-                    let eTag=eTag
-                        # we synchronize the eTags here because the order matters
-                        # for the final call to completeMultipartUpload
-                        put!(() -> push!(eTags, eTag), sync, n)
-                    end
-                end
+            Threads.@spawn begin
+                _n = $n
+                parteTag, wb = uploadPart(x, url, $part, _n, uploadState; credentials, kw...)
+                Threads.atomic_add!(wbytes, wb)
+                # we synchronize the eTags here because the order matters
+                # for the final call to completeMultipartUpload
+                put!(() -> push!(eTags, parteTag), sync, _n)
             end
         end
         eof(body) && break
@@ -97,5 +100,11 @@ function putObjectImpl(x::AbstractStore, key::String, in::RequestBodyType;
         close(body)
     end
     eTag = completeMultipartUpload(x, url, eTags, uploadState; credentials, kw...)
-    return Object(x, credentials, key, N, eTag)
+    obj = Object(x, credentials, key, N, eTag)
+@label done
+    end_time = time()
+    bytes = wbytes[]
+    gbits_per_second = bytes == 0 ? 0 : (((8 * bytes) / 1e9) / (end_time - start_time))
+    lograte && @info "CloudStore.put complete with bandwidth: $(gbits_per_second) Gbps"
+    return obj
 end

src/s3.jl

@@ -48,7 +48,7 @@ end
 
 function API.uploadPart(x::Bucket, url, part, partNumber, uploadId; kw...)
     resp = AWS.put(url, [], part; query=Dict("partNumber" => string(partNumber), "uploadId" => uploadId), service="s3", kw...)
-    return HTTP.header(resp, "ETag")
+    return (HTTP.header(resp, "ETag"), Base.get(resp.request.context, :nbytes_written, 0))
 end
 
 function API.completeMultipartUpload(x::Bucket, url, eTags, uploadId; kw...)