Fix after review

Author: ianton-ru

src/Disks/ObjectStorages/IObjectStorage.cpp

@@ -8,6 +8,10 @@
 #include <Common/Exception.h>
 #include <Common/ObjectStorageKeyGenerator.h>
 
+#include <Poco/JSON/Object.h>
+#include <Poco/JSON/Parser.h>
+#include <Poco/JSON/JSONException.h>
+
 
 namespace DB
 {
@@ -107,4 +111,36 @@ void RelativePathWithMetadata::loadMetadata(ObjectStoragePtr object_storage)
     }
 }
 
+RelativePathWithMetadata::CommandInTaskResponse::CommandInTaskResponse(const std::string & task)
+{
+    Poco::JSON::Parser parser;
+    try
+    {
+        auto json = parser.parse(task).extract<Poco::JSON::Object::Ptr>();
+        if (!json)
+            return;
+
+        successfully_parsed = true;
+
+        if (json->has("retry_after_us"))
+            retry_after_us = json->getValue<size_t>("retry_after_us");
+    }
+    catch (const Poco::JSON::JSONException &)
+    { /// Not a JSON
+        return;
+    }
+}
+
+std::string RelativePathWithMetadata::CommandInTaskResponse::to_string() const
+{
+    Poco::JSON::Object json;
+    if (retry_after_us.has_value())
+        json.set("retry_after_us", retry_after_us.value());
+
+    std::ostringstream oss;
+    oss.exceptions(std::ios::failbit);
+    Poco::JSON::Stringifier::stringify(json, oss);
+    return oss.str();
+}
+
 }

src/Disks/ObjectStorages/IObjectStorage.h

@@ -66,15 +66,37 @@ struct ObjectMetadata
 
 struct RelativePathWithMetadata
 {
+    class CommandInTaskResponse
+    {
+    public:
+        CommandInTaskResponse() {}
+        CommandInTaskResponse(const std::string & task);
+
+        bool is_parsed() const { return successfully_parsed; }
+        void set_retry_after_us(Poco::Timestamp::TimeDiff time_us) { retry_after_us = time_us; }
+
+        std::string to_string() const;
+
+        std::optional<Poco::Timestamp::TimeDiff> get_retry_after_us() const { return retry_after_us; }
+
+    private:
+        bool successfully_parsed = false;
+        std::optional<Poco::Timestamp::TimeDiff> retry_after_us;
+    };
+
     String relative_path;
     std::optional<ObjectMetadata> metadata;
+    CommandInTaskResponse command;
 
     RelativePathWithMetadata() = default;
 
-    explicit RelativePathWithMetadata(String relative_path_, std::optional<ObjectMetadata> metadata_ = std::nullopt)
-        : relative_path(std::move(relative_path_))
-        , metadata(std::move(metadata_))
-    {}
+    explicit RelativePathWithMetadata(const String & task_string, std::optional<ObjectMetadata> metadata_ = std::nullopt)
+        : metadata(std::move(metadata_))
+        , command(task_string)
+    {
+        if (!command.is_parsed())
+            relative_path = task_string;
+    }
 
     virtual ~RelativePathWithMetadata() = default;
 
@@ -85,6 +107,8 @@ struct RelativePathWithMetadata
     virtual size_t fileSizeInArchive() const { throw Exception(ErrorCodes::LOGICAL_ERROR, "Not an archive"); }
 
     void loadMetadata(ObjectStoragePtr object_storage);
+
+    const CommandInTaskResponse & getCommand() const { return command; }
 };
 
 struct ObjectKeyWithMetadata

src/Storages/ObjectStorage/StorageObjectStorageCluster.cpp

@@ -35,6 +35,7 @@ namespace ErrorCodes
     extern const int LOGICAL_ERROR;
     extern const int UNKNOWN_FUNCTION;
     extern const int NOT_IMPLEMENTED;
+    extern const int INVALID_SETTING_VALUE;
 }
 
 
@@ -387,10 +388,22 @@ RemoteQueryExecutor::Extension StorageObjectStorageCluster::getTaskIteratorExten
         }
     }
 
+    uint64_t lock_object_storage_task_distribution_ms = local_context->getSettingsRef()[Setting::lock_object_storage_task_distribution_ms];
+
+    /// Check value to avoid negative result after conversion in microseconds.
+    /// Poco::Timestamp::TimeDiff is signed int 64.
+    static const uint64_t lock_object_storage_task_distribution_ms_max = 0x0020000000000000ULL;
+    if (lock_object_storage_task_distribution_ms > lock_object_storage_task_distribution_ms_max)
+        throw Exception(ErrorCodes::INVALID_SETTING_VALUE,
+            "Value lock_object_storage_task_distribution_ms is too big: {}, allowed maximum is {}",
+            lock_object_storage_task_distribution_ms,
+            lock_object_storage_task_distribution_ms_max
+        );
+
     auto task_distributor = std::make_shared<StorageObjectStorageStableTaskDistributor>(
         iterator,
         ids_of_hosts,
-        local_context->getSettingsRef()[Setting::lock_object_storage_task_distribution_ms]);
+        lock_object_storage_task_distribution_ms);
 
     auto callback = std::make_shared<TaskIterator>(
         [task_distributor](size_t number_of_current_replica) mutable -> String {

src/Storages/ObjectStorage/StorageObjectStorageSource.cpp

@@ -438,22 +438,26 @@ StorageObjectStorageSource::ReaderHolder StorageObjectStorageSource::createReade
         not_a_path = false;
         object_info = file_iterator->next(processor);
 
-        if (!object_info || object_info->getPath().empty())
+        if (!object_info)
             return {};
 
-        StorageObjectStorageStableTaskDistributor::CommandInTaskResponse command(object_info->getPath());
-        if (command.is_parsed())
+        if (object_info->getCommand().is_parsed())
         {
-            auto retry_after_us = command.get_retry_after_us();
+            auto retry_after_us = object_info->getCommand().get_retry_after_us();
             if (retry_after_us.has_value())
             {
                 not_a_path = true;
                 /// TODO: Make asyncronous waiting without sleep in thread
-                sleepForMicroseconds(std::min(100000ul, retry_after_us.value()));
+                /// Now this sleep is on executor node in worker thread
+                /// Does not block query initiator
+                sleepForMicroseconds(std::min(Poco::Timestamp::TimeDiff(100000ul), retry_after_us.value()));
                 continue;
             }
         }
 
+        if (object_info->getPath().empty())
+            return {};
+
         object_info->loadMetadata(object_storage);
     }
     while (not_a_path || (query_settings.skip_empty_files && object_info->metadata->size_bytes == 0));

src/Storages/ObjectStorage/StorageObjectStorageStableTaskDistributor.cpp

@@ -3,10 +3,6 @@
 #include <consistent_hashing.h>
 #include <optional>
 
-#include <Poco/JSON/Object.h>
-#include <Poco/JSON/Parser.h>
-#include <Poco/JSON/JSONException.h>
-
 namespace DB
 {
 
@@ -169,7 +165,7 @@ std::optional<String> StorageObjectStorageStableTaskDistributor::getAnyUnprocess
     /// Limit time of node activity to keep task in queue
     Poco::Timestamp activity_limit;
     Poco::Timestamp oldest_activity;
-    if (lock_object_storage_task_distribution_us)
+    if (lock_object_storage_task_distribution_us > 0)
         activity_limit -= lock_object_storage_task_distribution_us;
 
     std::lock_guard lock(mutex);
@@ -181,7 +177,7 @@ std::optional<String> StorageObjectStorageStableTaskDistributor::getAnyUnprocess
         while (it != unprocessed_files.end())
         {
             auto last_activity = last_node_activity.find(it->second);
-            if (!lock_object_storage_task_distribution_us
+            if (lock_object_storage_task_distribution_us <= 0
                 || last_activity == last_node_activity.end()
                 || activity_limit > last_activity->second)
             {
@@ -211,7 +207,7 @@ std::optional<String> StorageObjectStorageStableTaskDistributor::getAnyUnprocess
 
         /// All unprocessed files owned by alive replicas with recenlty activity
         /// Need to retry after (oldest_activity - activity_limit) microseconds
-        CommandInTaskResponse response;
+        RelativePathWithMetadata::CommandInTaskResponse response;
         response.set_retry_after_us(oldest_activity - activity_limit);
         return response.to_string();
     }
@@ -226,36 +222,4 @@ void StorageObjectStorageStableTaskDistributor::saveLastNodeActivity(size_t numb
     last_node_activity[number_of_current_replica] = now;
 }
 
-StorageObjectStorageStableTaskDistributor::CommandInTaskResponse::CommandInTaskResponse(const std::string & task)
-{
-    Poco::JSON::Parser parser;
-    try
-    {
-        auto json = parser.parse(task).extract<Poco::JSON::Object::Ptr>();
-        if (!json)
-            return;
-
-        successfully_parsed = true;
-
-        if (json->has("retry_after_us"))
-            retry_after_us = json->getValue<size_t>("retry_after_us");
-    }
-    catch (const Poco::JSON::JSONException &)
-    { /// Not a JSON
-        return;
-    }
-}
-
-std::string StorageObjectStorageStableTaskDistributor::CommandInTaskResponse::to_string() const
-{
-    Poco::JSON::Object json;
-    if (retry_after_us.has_value())
-        json.set("retry_after_us", retry_after_us.value());
-
-    std::ostringstream oss;
-    oss.exceptions(std::ios::failbit);
-    Poco::JSON::Stringifier::stringify(json, oss);
-    return oss.str();
-}
-
 }

src/Storages/ObjectStorage/StorageObjectStorageStableTaskDistributor.h

@@ -20,24 +20,6 @@ namespace DB
 class StorageObjectStorageStableTaskDistributor
 {
 public:
-    class CommandInTaskResponse
-    {
-    public:
-        CommandInTaskResponse() {}
-        CommandInTaskResponse(const std::string & task);
-
-        bool is_parsed() const { return successfully_parsed; }
-        void set_retry_after_us(uint64_t time_us) { retry_after_us = time_us; }
-
-        std::string to_string() const;
-
-        std::optional<uint64_t> get_retry_after_us() const { return retry_after_us; }
-
-    private:
-        bool successfully_parsed = false;
-        std::optional<uint64_t> retry_after_us;
-    };
-
     StorageObjectStorageStableTaskDistributor(
         std::shared_ptr<IObjectIterator> iterator_,
         std::vector<std::string> ids_of_nodes_,

tests/integration/test_s3_cache_locality/test.py

@@ -75,7 +75,8 @@ def started_cluster():
         cluster.shutdown()
 
 
-def check_s3_gets(cluster, node, expected_result, cluster_first, cluster_second, enable_filesystem_cache, lock=False):
+def check_s3_gets(cluster, node, expected_result, cluster_first, cluster_second, enable_filesystem_cache,
+                  lock_object_storage_task_distribution_ms):
     for host in list(cluster.instances.values()):
         host.query("SYSTEM DROP FILESYSTEM CACHE 'raw_s3_cache'", ignore_error=True)
 
@@ -84,8 +85,8 @@ def check_s3_gets(cluster, node, expected_result, cluster_first, cluster_second,
         "filesystem_cache_name": "'raw_s3_cache'",
         }
 
-    if lock:
-        settings["lock_object_storage_task_distribution_ms"] = 30000
+    if lock_object_storage_task_distribution_ms > 0:
+        settings["lock_object_storage_task_distribution_ms"] = lock_object_storage_task_distribution_ms
 
     query_id_first = str(uuid.uuid4())
     result_first = node.query(
@@ -133,20 +134,21 @@ def check_s3_gets(cluster, node, expected_result, cluster_first, cluster_second,
     return int(s3_get_first), int(s3_get_second)
 
 
-def check_s3_gets_repeat(cluster, node, expected_result, cluster_first, cluster_second, enable_filesystem_cache, lock=False):
+def check_s3_gets_repeat(cluster, node, expected_result, cluster_first, cluster_second, enable_filesystem_cache,
+                         lock_object_storage_task_distribution_ms):
     # Repeat test several times to get average result
-    iterations = 1 if lock else 10
+    iterations = 1 if lock_object_storage_task_distribution_ms > 0 else 10
     s3_get_first_sum = 0
     s3_get_second_sum = 0
     for _ in range(iterations):
-        (s3_get_first, s3_get_second) = check_s3_gets(cluster, node, expected_result, cluster_first, cluster_second, enable_filesystem_cache, lock)
+        (s3_get_first, s3_get_second) = check_s3_gets(cluster, node, expected_result, cluster_first, cluster_second, enable_filesystem_cache, lock_object_storage_task_distribution_ms)
         s3_get_first_sum += s3_get_first
         s3_get_second_sum += s3_get_second
     return s3_get_first_sum, s3_get_second_sum
 
 
-@pytest.mark.parametrize("lock", [False, True])
-def test_cache_locality(started_cluster, lock):
+@pytest.mark.parametrize("lock_object_storage_task_distribution_ms ", [0, 30000])
+def test_cache_locality(started_cluster, lock_object_storage_task_distribution_ms):
     node = started_cluster.instances["clickhouse0"]
 
     expected_result = node.query(
@@ -158,36 +160,36 @@ def test_cache_locality(started_cluster, lock):
     )
 
     # Algorithm does not give 100% guarantee, so add 10% on dispersion
-    dispersion = 0.0 if lock else 0.1
+    dispersion = 0.0 if lock_object_storage_task_distribution_ms > 0 else 0.1
 
     # No cache
-    (s3_get_first, s3_get_second) = check_s3_gets_repeat(started_cluster, node, expected_result, 'cluster_12345', 'cluster_12345', 0, lock)
+    (s3_get_first, s3_get_second) = check_s3_gets_repeat(started_cluster, node, expected_result, 'cluster_12345', 'cluster_12345', 0, lock_object_storage_task_distribution_ms)
     assert s3_get_second == s3_get_first
 
     # With cache
-    (s3_get_first, s3_get_second) = check_s3_gets_repeat(started_cluster, node, expected_result, 'cluster_12345', 'cluster_12345', 1, lock)
+    (s3_get_first, s3_get_second) = check_s3_gets_repeat(started_cluster, node, expected_result, 'cluster_12345', 'cluster_12345', 1, lock_object_storage_task_distribution_ms)
     assert s3_get_second <= s3_get_first * dispersion
 
     # Different nodes order
-    (s3_get_first, s3_get_second) = check_s3_gets_repeat(started_cluster, node, expected_result, 'cluster_12345', 'cluster_34512', 1, lock)
+    (s3_get_first, s3_get_second) = check_s3_gets_repeat(started_cluster, node, expected_result, 'cluster_12345', 'cluster_34512', 1, lock_object_storage_task_distribution_ms)
     assert s3_get_second <= s3_get_first * dispersion
 
     # No last node
-    (s3_get_first, s3_get_second) = check_s3_gets_repeat(started_cluster, node, expected_result, 'cluster_12345', 'cluster_1234', 1, lock)
+    (s3_get_first, s3_get_second) = check_s3_gets_repeat(started_cluster, node, expected_result, 'cluster_12345', 'cluster_1234', 1, lock_object_storage_task_distribution_ms)
     assert s3_get_second <= s3_get_first * (0.211 + dispersion) # actual value - 24 for 100 files, 211 for 1000
 
     # No first node
-    (s3_get_first, s3_get_second) = check_s3_gets_repeat(started_cluster, node, expected_result, 'cluster_12345', 'cluster_2345', 1, lock)
+    (s3_get_first, s3_get_second) = check_s3_gets_repeat(started_cluster, node, expected_result, 'cluster_12345', 'cluster_2345', 1, lock_object_storage_task_distribution_ms)
     assert s3_get_second <= s3_get_first * (0.189 + dispersion) # actual value - 12 for 100 files, 189 for 1000
 
     # No first node, different nodes order
-    (s3_get_first, s3_get_second) = check_s3_gets_repeat(started_cluster, node, expected_result, 'cluster_12345', 'cluster_4523', 1, lock)
+    (s3_get_first, s3_get_second) = check_s3_gets_repeat(started_cluster, node, expected_result, 'cluster_12345', 'cluster_4523', 1, lock_object_storage_task_distribution_ms)
     assert s3_get_second <= s3_get_first * (0.189 + dispersion)
 
     # Add new node, different nodes order
-    (s3_get_first, s3_get_second) = check_s3_gets_repeat(started_cluster, node, expected_result, 'cluster_4523', 'cluster_12345', 1, lock)
+    (s3_get_first, s3_get_second) = check_s3_gets_repeat(started_cluster, node, expected_result, 'cluster_4523', 'cluster_12345', 1, lock_object_storage_task_distribution_ms)
     assert s3_get_second <= s3_get_first * (0.189 + dispersion)
 
     # New node and old node, different nodes order
-    (s3_get_first, s3_get_second) = check_s3_gets_repeat(started_cluster, node, expected_result, 'cluster_1234', 'cluster_4523', 1, lock)
+    (s3_get_first, s3_get_second) = check_s3_gets_repeat(started_cluster, node, expected_result, 'cluster_1234', 'cluster_4523', 1, lock_object_storage_task_distribution_ms)
     assert s3_get_second <= s3_get_first * (0.400 + dispersion) # actual value - 36 for 100 files, 400 for 1000