diff --git a/cpp/distributed/client.cpp b/cpp/distributed/client.cpp index d6aad928a..49c4800fe 100644 --- a/cpp/distributed/client.cpp +++ b/cpp/distributed/client.cpp @@ -207,11 +207,15 @@ static httplib::Result oneShotDownload( const Url& proxyUrl, size_t startByte, //inclusive size_t endByte, //inclusive - std::function f + httplib::ResponseHandler responseHandler, + httplib::ContentReceiver contentReceiver ) { httplib::Headers headers; if(startByte > 0) { - headers.insert(std::make_pair("Range", Global::uint64ToString(startByte) + "-" + Global::uint64ToString(endByte))); + //Per RFC 7233 the Range header requires an explicit byte-range unit. Servers that + //receive a Range header without one treat it as malformed and ignore it, replying + //with 200 and the full content, which corrupts resumed downloads. + headers.insert(std::make_pair("Range", "bytes=" + Global::uint64ToString(startByte) + "-" + Global::uint64ToString(endByte))); } if(!url.isSSL) { @@ -224,7 +228,7 @@ static httplib::Result oneShotDownload( } //Avoid automatically decompressing .bin.gz files that get sent to us with "content-encoding: gzip" httpClient->set_decompress(false); - return httpClient->Get(url.path.c_str(),headers,f); + return httpClient->Get(url.path.c_str(),headers,responseHandler,contentReceiver); } else { std::unique_ptr httpsClient = std::make_unique(url.host, url.port); @@ -238,7 +242,7 @@ static httplib::Result oneShotDownload( httpsClient->enable_server_certificate_verification(true); //Avoid automatically decompressing .bin.gz files that get sent to us with "content-encoding: gzip" httpsClient->set_decompress(false); - httplib::Result response = httpsClient->Get(url.path.c_str(),headers,f); + httplib::Result response = httpsClient->Get(url.path.c_str(),headers,responseHandler,contentReceiver); if(response == nullptr) { auto result = httpsClient->get_openssl_verify_result(); if(result) { @@ -1012,16 +1016,83 @@ bool Connection::actuallyDownloadModel( double lastTime = timer.getSeconds(); size_t totalDataSize = 0; + bool serverIgnoresRange = false; + + //Throw away all partial data so far and begin again from byte 0. Used whenever we cannot + //trust that the bytes we have so far are a prefix of the true file, e.g. when a retry got + //a response that was not a valid continuation of the ranged request we made. + auto restartFromBeginning = [&]() { + out.close(); + FileUtils::open(out,tmpPath,ios::binary); + totalDataSize = 0; + }; auto fInner = [&](int& innerLoopFailMode) { - if(totalDataSize >= modelInfo.bytes) + if(totalDataSize > modelInfo.bytes) { + restartFromBeginning(); + throw StringError( + "Model download received more bytes than the model contains, discarding partial download and restarting from the beginning" + ); + } + //If we already learned that the server will not honor Range headers, resuming is not + //possible at all, so any partial data must be thrown away and downloaded in one piece. + if(serverIgnoresRange && totalDataSize > 0) + restartFromBeginning(); + if(totalDataSize == modelInfo.bytes) return; const size_t oldTotalDataSize = totalDataSize; const size_t startByte = oldTotalDataSize; const size_t endByte = modelInfo.bytes-1; + const bool requestResume = startByte > 0; + if(requestResume) + logger->write("Resuming download of model at byte " + Global::uint64ToString(startByte) + " of " + Global::uint64ToString(modelInfo.bytes) + ": " + urlToActuallyUse.originalString); const Url proxyToUse = mirrorUseProxy ? proxyUrl : Url(); + + //Validate response headers before the content receiver appends any bytes. In particular, + //cpp-httplib returns a null Result if the content receiver cancels after an overshoot, so + //waiting until after the body has been streamed would lose the 200 status that tells us + //the server ignored Range. + string responseHeaderError; + bool disableRangeResume = false; + httplib::ResponseHandler responseHandler = [&](const httplib::Response& headerResponse) { + if(headerResponse.status != 200 && headerResponse.status != 206) { + ostringstream outs; + outs << "Server gave response status code " << headerResponse.status << " instead of 200 OK or 206 Partial Content"; + for(const auto& header : headerResponse.headers) + outs << "\nHeader: " << header.first << ": " << header.second; + responseHeaderError = outs.str(); + return false; + } + + if(requestResume && headerResponse.status == 200) { + disableRangeResume = true; + responseHeaderError = + "Server ignored the Range header on a resumed download and sent the full file, " + "discarding partial download and restarting from the beginning"; + return false; + } + + if(requestResume && headerResponse.status == 206) { + string contentRange; + for(const auto& header : headerResponse.headers) { + if(Global::toLower(header.first) == "content-range") + contentRange = header.second; + } + const string expectedPrefix = string("bytes ") + Global::uint64ToString(startByte) + "-"; + if(contentRange == "" || Global::toLower(contentRange).rfind(expectedPrefix,0) != 0) { + disableRangeResume = true; + responseHeaderError = + "Server returned 206 Partial Content with " + + (contentRange == "" ? string("no Content-Range header") : string("unexpected Content-Range \"") + contentRange + "\"") + + ", discarding partial download and restarting from the beginning"; + return false; + } + } + return true; + }; + httplib::Result response = oneShotDownload( - logger, urlToActuallyUse, caCertsFile, proxyToUse, startByte, endByte, + logger, urlToActuallyUse, caCertsFile, proxyToUse, startByte, endByte, responseHandler, [&out,&totalDataSize,&shouldStop,this,&timer,&lastTime,&urlToActuallyUse,&modelInfo](const char* data, size_t data_length) { out.write(data, data_length); totalDataSize += data_length; @@ -1043,8 +1114,26 @@ bool Connection::actuallyDownloadModel( if(shouldStop()) throw StringError("Stopping because shouldStop is true"); - if(totalDataSize > oldTotalDataSize) - innerLoopFailMode = LOOP_PARTIAL_SUCCESS; + if(responseHeaderError != "") { + if(disableRangeResume) { + serverIgnoresRange = true; + restartFromBeginning(); + } + throw StringError(responseHeaderError); + } + + if(totalDataSize > modelInfo.bytes) { + restartFromBeginning(); + throw StringError( + "Model download received more bytes than the model contains, discarding partial download and restarting from the beginning" + ); + } + + if(totalDataSize > oldTotalDataSize) { + //Only retained bytes count as partial success. If Range is unavailable, the next retry + //must discard this incomplete full response, so it must consume the normal retry budget. + innerLoopFailMode = serverIgnoresRange ? LOOP_RETRYABLE_FAIL : LOOP_PARTIAL_SUCCESS; + } if(response == nullptr) throw StringError("No response from server");