From fcc6c2024f04b2fbfcfc25c1633f5747b032fd94 Mon Sep 17 00:00:00 2001
From: JinmingHu <jinmhu@microsoft.com>
Date: Tue, 6 Apr 2021 11:19:04 +0800
Subject: [PATCH] server timeout support (#2010)

---
 .../azure-storage-blobs/CMakeLists.txt        |  1 +
 .../test/ut/storage_timeout_test.cpp          | 72 +++++++++++++++++++
 sdk/storage/azure-storage-common/CHANGELOG.md |  4 ++
 .../src/storage_per_retry_policy.cpp          | 17 +++++
 4 files changed, 94 insertions(+)
 create mode 100644 sdk/storage/azure-storage-blobs/test/ut/storage_timeout_test.cpp

diff --git a/sdk/storage/azure-storage-blobs/CMakeLists.txt b/sdk/storage/azure-storage-blobs/CMakeLists.txt
index 5fddf6c33..1ad666cd5 100644
--- a/sdk/storage/azure-storage-blobs/CMakeLists.txt
+++ b/sdk/storage/azure-storage-blobs/CMakeLists.txt
@@ -100,6 +100,7 @@ if(BUILD_TESTING)
         test/ut/page_blob_client_test.cpp
         test/ut/page_blob_client_test.hpp
         test/ut/storage_retry_policy_test.cpp
+        test/ut/storage_timeout_test.cpp
   )
 
   target_link_libraries(azure-storage-test PRIVATE azure-storage-blobs)
diff --git a/sdk/storage/azure-storage-blobs/test/ut/storage_timeout_test.cpp b/sdk/storage/azure-storage-blobs/test/ut/storage_timeout_test.cpp
new file mode 100644
index 000000000..6e1128c78
--- /dev/null
+++ b/sdk/storage/azure-storage-blobs/test/ut/storage_timeout_test.cpp
@@ -0,0 +1,72 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// SPDX-License-Identifier: MIT
+
+#include <chrono>
+#include <memory>
+#include <vector>
+
+#include <azure/storage/blobs.hpp>
+
+#include "test_base.hpp"
+
+namespace Azure { namespace Storage { namespace Test {
+
+  class PeekHttpRequestPolicy : public Core::Http::Policies::HttpPolicy {
+  public:
+    PeekHttpRequestPolicy(std::function<void(const Core::Http::Request&)> callback)
+        : m_callback(std::move(callback))
+    {
+    }
+
+    std::unique_ptr<Core::Http::RawResponse> Send(
+        Core::Http::Request& request,
+        Core::Http::Policies::NextHttpPolicy nextHttpPolicy,
+        Core::Context const& context) const override
+    {
+      m_callback(request);
+      return nextHttpPolicy.Send(request, context);
+    }
+
+    std::unique_ptr<HttpPolicy> Clone() const override
+    {
+      return std::make_unique<PeekHttpRequestPolicy>(*this);
+    }
+
+  private:
+    std::function<void(const Core::Http::Request&)> m_callback;
+  };
+
+  TEST(StoragetimeoutTest, Basic)
+  {
+    Azure::Nullable<int64_t> timeout;
+    auto callback = [&timeout](const Core::Http::Request& request) {
+      auto queryParameteres = request.GetUrl().GetQueryParameters();
+      auto ite = queryParameteres.find("timeout");
+      if (ite == queryParameteres.end())
+      {
+        timeout.Reset();
+      }
+      else
+      {
+        timeout = std::stoll(ite->second);
+      }
+    };
+
+    auto peekPolicyPtr = std::make_unique<PeekHttpRequestPolicy>(callback);
+    Blobs::BlobClientOptions clientOptions;
+    clientOptions.PerRetryPolicies.emplace_back(std::move(peekPolicyPtr));
+    auto containerClient = Azure::Storage::Blobs::BlobContainerClient::CreateFromConnectionString(
+        StandardStorageConnectionString(), LowercaseRandomString(), clientOptions);
+    containerClient.DeleteIfExists();
+    EXPECT_FALSE(timeout.HasValue());
+
+    Azure::Core::Context context;
+    context = context.WithDeadline(std::chrono::system_clock::now() + std::chrono::seconds(300));
+    containerClient.DeleteIfExists(Storage::Blobs::DeleteBlobContainerOptions(), context);
+
+    ASSERT_TRUE(timeout.HasValue());
+    EXPECT_GE(timeout.Value(), 299);
+    EXPECT_LE(timeout.Value(), 301);
+  }
+
+}}} // namespace Azure::Storage::Test
diff --git a/sdk/storage/azure-storage-common/CHANGELOG.md b/sdk/storage/azure-storage-common/CHANGELOG.md
index 3098fa615..54602e7bb 100644
--- a/sdk/storage/azure-storage-common/CHANGELOG.md
+++ b/sdk/storage/azure-storage-common/CHANGELOG.md
@@ -2,6 +2,10 @@
 
 ## 12.0.0-beta.10 (Unreleased)
 
+### New Features
+
+- Added server timeout support.
+
 ### Breaking Changes
 
 - Removed `Azure::Storage::Common::PackageVersion`.
diff --git a/sdk/storage/azure-storage-common/src/storage_per_retry_policy.cpp b/sdk/storage/azure-storage-common/src/storage_per_retry_policy.cpp
index edda113c6..4b0207ac9 100644
--- a/sdk/storage/azure-storage-common/src/storage_per_retry_policy.cpp
+++ b/sdk/storage/azure-storage-common/src/storage_per_retry_policy.cpp
@@ -6,6 +6,7 @@
 #include <azure/core/datetime.hpp>
 #include <azure/core/platform.hpp>
 
+#include <algorithm>
 #include <chrono>
 
 namespace Azure { namespace Storage { namespace _internal {
@@ -28,6 +29,22 @@ namespace Azure { namespace Storage { namespace _internal {
               .ToString(Azure::DateTime::DateFormat::Rfc1123));
     }
 
+    const char* HttpHeaderTimeout = "timeout";
+    auto cancelTimepoint = ctx.GetDeadline();
+    if (cancelTimepoint == Azure::DateTime::max())
+    {
+      request.GetUrl().RemoveQueryParameter(HttpHeaderTimeout);
+    }
+    else
+    {
+      auto currentTimepoint = std::chrono::system_clock::now();
+      int64_t numSeconds
+          = std::chrono::duration_cast<std::chrono::seconds>(
+                std::chrono::system_clock::time_point(cancelTimepoint) - currentTimepoint)
+                .count();
+      request.GetUrl().AppendQueryParameter(
+          HttpHeaderTimeout, std::to_string(std::max(numSeconds, int64_t(1))));
+    }
     return nextHttpPolicy.Send(request, ctx);
   }