Enhance K2 TL lib (#1109)

* enhance TL for job workers * add some general TL types * add TL types and functions for K2 confdata
VKCOM · Sep 26, 2024 · 38b4837 · 38b4837
1 parent ee1122d
commit 38b4837
Show file tree

Hide file tree

Showing 8 changed files with 330 additions and 33 deletions.
diff --git a/runtime-core/memory-resource/resource_allocator.h b/runtime-core/memory-resource/resource_allocator.h
@@ -9,9 +9,9 @@
 #include <queue>
 #include <unordered_map>
 #include <unordered_set>
+#include <vector>
 
 #include "common/wrappers/likely.h"
-
 #include "runtime-core/utils/kphp-assert-core.h"
 
 namespace memory_resource {
@@ -79,6 +79,9 @@ using deque = std::deque<T, resource_allocator<T, Resource>>;
 
 template<class T, class Resource>
 using queue = std::queue<T, std::deque<T, resource_allocator<T, Resource>>>;
+
+template<class T, class Resource>
+using vector = std::vector<T, resource_allocator<T, Resource>>;
 } // namespace stl
 
 } // namespace memory_resource
diff --git a/runtime-light/stdlib/crypto/crypto-functions.cpp b/runtime-light/stdlib/crypto/crypto-functions.cpp
@@ -12,7 +12,7 @@
 
 namespace {
 
-const string CRYPTO_COMPONENT = string("crypto");
+constexpr const char *CRYPTO_COMPONENT_NAME = "crypto";
 
 } // namespace
 
@@ -31,7 +31,7 @@ task_t<Optional<string>> f$openssl_random_pseudo_bytes(int64_t length) noexcept
   string request_buf;
   request_buf.append(buffer.data(), buffer.size());
 
-  auto query = f$component_client_send_request(CRYPTO_COMPONENT, string{buffer.data(), static_cast<string::size_type>(buffer.size())});
+  auto query = f$component_client_send_request(string{CRYPTO_COMPONENT_NAME}, string{buffer.data(), static_cast<string::size_type>(buffer.size())});
   string resp = co_await f$component_client_fetch_response(co_await query);
 
   buffer.clean();
@@ -55,7 +55,7 @@ task_t<Optional<array<mixed>>> f$openssl_x509_parse(const string &data, bool sho
   string request_buf;
   request_buf.append(buffer.data(), buffer.size());
 
-  auto query = f$component_client_send_request(CRYPTO_COMPONENT, request_buf);
+  auto query = f$component_client_send_request(string{CRYPTO_COMPONENT_NAME}, request_buf);
   string resp_from_platform = co_await f$component_client_fetch_response(co_await query);
 
   buffer.clean();
@@ -75,7 +75,7 @@ task_t<bool> f$openssl_sign(const string &data, string &signature, const string
   tl::TLBuffer buffer;
   request.store(buffer);
 
-  auto query = f$component_client_send_request(CRYPTO_COMPONENT, string(buffer.data(), buffer.size()));
+  auto query = f$component_client_send_request(string{CRYPTO_COMPONENT_NAME}, string(buffer.data(), buffer.size()));
   string resp_from_platform = co_await f$component_client_fetch_response(co_await query);
 
   buffer.clean();
@@ -98,7 +98,7 @@ task_t<int64_t> f$openssl_verify(const string &data, const string &signature, co
   tl::TLBuffer buffer;
   request.store(buffer);
 
-  auto query = f$component_client_send_request(CRYPTO_COMPONENT, string(buffer.data(), buffer.size()));
+  auto query = f$component_client_send_request(string{CRYPTO_COMPONENT_NAME}, string(buffer.data(), buffer.size()));
   string resp_from_platform = co_await f$component_client_fetch_response(co_await query);
 
   buffer.clean();

diff --git a/runtime-light/tl/tl-core.cpp b/runtime-light/tl/tl-core.cpp
@@ -5,6 +5,7 @@
 #include "runtime-light/tl/tl-core.h"
 
 namespace tl {
+
 void TLBuffer::store_string(std::string_view str) noexcept {
   const char *str_buf{str.data()};
   size_t str_len{str.size()};
@@ -94,4 +95,5 @@ std::string_view TLBuffer::fetch_string() noexcept {
   adjust(total_len_with_padding - size_len);
   return response;
 }
+
 } // namespace tl
diff --git a/runtime-light/tl/tl-core.h b/runtime-light/tl/tl-core.h
@@ -4,6 +4,7 @@
 
 #pragma once
 
+#include <concepts>
 #include <cstddef>
 #include <optional>
 #include <string_view>
@@ -26,7 +27,7 @@ inline constexpr uint64_t MEDIUM_STRING_MAX_LEN = (static_cast<uint64_t>(1) << 2
 inline constexpr uint8_t LARGE_STRING_MAGIC = 0xff;
 inline constexpr uint8_t MEDIUM_STRING_MAGIC = 0xfe;
 
-class TLBuffer : private vk::not_copyable {
+class TLBuffer final : private vk::not_copyable {
   string_buffer m_buffer;
   size_t m_pos{0};
   size_t m_remaining{0};
@@ -57,7 +58,7 @@ class TLBuffer : private vk::not_copyable {
   }
 
   void reset(size_t pos) noexcept {
-    php_assert(pos >= 0 && pos <= size());
+    php_assert(pos <= size());
     m_pos = pos;
     m_remaining = size() - m_pos;
   }
@@ -97,4 +98,14 @@ class TLBuffer : private vk::not_copyable {
   std::string_view fetch_string() noexcept;
 };
 
+template<typename T>
+concept tl_serializable = std::default_initializable<T> && requires(T t, TLBuffer tlb) {
+  { t.store(tlb) } noexcept -> std::same_as<void>;
+};
+
+template<typename T>
+concept tl_deserializable = std::default_initializable<T> && requires(T t, TLBuffer tlb) {
+  { t.fetch(tlb) } noexcept -> std::convertible_to<bool>;
+};
+
 } // namespace tl
diff --git a/runtime-light/tl/tl-functions.cpp b/runtime-light/tl/tl-functions.cpp
@@ -4,6 +4,7 @@
 
 #include "runtime-light/tl/tl-functions.h"
 
+#include <cstddef>
 #include <cstdint>
 #include <string_view>
 
@@ -12,27 +13,31 @@
 
 namespace {
 
-// magic + flags + image_id + job_id + minimum string size length
-constexpr auto K2_INVOKE_JW_MIN_SIZE = sizeof(uint32_t) + sizeof(uint32_t) + sizeof(uint64_t) + sizeof(int64_t) + sizeof(uint64_t) + tl::SMALL_STRING_SIZE_LEN;
-
 constexpr auto K2_JOB_WORKER_IGNORE_ANSWER_FLAG = static_cast<uint32_t>(1U << 0U);
 
 } // namespace
 
 namespace tl {
 
 bool K2InvokeJobWorker::fetch(TLBuffer &tlb) noexcept {
-  if (tlb.size() < K2_INVOKE_JW_MIN_SIZE || *tlb.fetch_trivial<uint32_t>() != K2_INVOKE_JOB_WORKER_MAGIC) {
+  if (tlb.fetch_trivial<uint32_t>().value_or(TL_ZERO) != K2_INVOKE_JOB_WORKER_MAGIC) {
+    return false;
+  }
+
+  const auto opt_flags{tlb.fetch_trivial<uint32_t>()};
+  const auto opt_image_id{tlb.fetch_trivial<uint64_t>()};
+  const auto opt_job_id{tlb.fetch_trivial<int64_t>()};
+  const auto opt_timeout_ns{tlb.fetch_trivial<uint64_t>()};
+  if (!(opt_flags.has_value() && opt_image_id.has_value() && opt_job_id.has_value() && opt_timeout_ns.has_value())) {
     return false;
   }
+  const auto body_view{tlb.fetch_string()};
 
-  const auto flags{*tlb.fetch_trivial<uint32_t>()};
-  ignore_answer = static_cast<bool>(flags & K2_JOB_WORKER_IGNORE_ANSWER_FLAG);
-  image_id = *tlb.fetch_trivial<uint64_t>();
-  job_id = *tlb.fetch_trivial<int64_t>();
-  timeout_ns = *tlb.fetch_trivial<uint64_t>();
-  const std::string_view body_view{tlb.fetch_string()};
-  body = string{body_view.data(), static_cast<string::size_type>(body_view.size())};
+  ignore_answer = static_cast<bool>(*opt_flags & K2_JOB_WORKER_IGNORE_ANSWER_FLAG);
+  image_id = *opt_image_id;
+  job_id = *opt_job_id;
+  timeout_ns = *opt_timeout_ns;
+  body = {body_view.data(), static_cast<string::size_type>(body_view.size())};
   return true;
 }
 
@@ -72,4 +77,14 @@ void DigestVerify::store(TLBuffer &tlb) const noexcept {
   tlb.store_string(std::string_view{signature.c_str(), signature.size()});
 }
 
+void ConfdataGet::store(TLBuffer &tlb) const noexcept {
+  tlb.store_trivial<uint32_t>(CONFDATA_GET_MAGIC);
+  tlb.store_string({key.c_str(), static_cast<size_t>(key.size())});
+}
+
+void ConfdataGetWildcard::store(TLBuffer &tlb) const noexcept {
+  tlb.store_trivial<uint32_t>(CONFDATA_GET_WILDCARD_MAGIC);
+  tlb.store_string({wildcard.c_str(), static_cast<size_t>(wildcard.size())});
+}
+
 } // namespace tl
diff --git a/runtime-light/tl/tl-functions.h b/runtime-light/tl/tl-functions.h
@@ -36,7 +36,6 @@ inline constexpr uint32_t GET_PEM_CERT_INFO_MAGIC = 0xa50c'fd6c;
 inline constexpr uint32_t DIGEST_SIGN_MAGIC = 0xd345'f658;
 inline constexpr uint32_t DIGEST_VERIFY_MAGIC = 0x5760'bd0e;
 
-
 struct GetCryptosecurePseudorandomBytes final {
   int32_t size{};
 
@@ -67,4 +66,21 @@ struct DigestVerify final {
   void store(TLBuffer &tlb) const noexcept;
 };
 
+// ===== CONFDATA =====
+
+inline constexpr uint32_t CONFDATA_GET_MAGIC = 0xf0eb'cd89;
+inline constexpr uint32_t CONFDATA_GET_WILDCARD_MAGIC = 0x5759'bd9e;
+
+struct ConfdataGet final {
+  string key;
+
+  void store(TLBuffer &tlb) const noexcept;
+};
+
+struct ConfdataGetWildcard final {
+  string wildcard;
+
+  void store(TLBuffer &tlb) const noexcept;
+};
+
 } // namespace tl
diff --git a/runtime-light/tl/tl-types.cpp b/runtime-light/tl/tl-types.cpp
@@ -4,6 +4,7 @@
 
 #include "runtime-light/tl/tl-types.h"
 
+#include <cstddef>
 #include <cstdint>
 #include <string_view>
 #include <tuple>
@@ -13,27 +14,28 @@
 
 namespace {
 
-// magic + flags + job_id + minimum string size length
-constexpr auto K2_JOB_WORKER_RESPONSE_MIN_SIZE = sizeof(uint32_t) + sizeof(uint32_t) + sizeof(int64_t) + tl::SMALL_STRING_SIZE_LEN;
+enum CertInfoItem : uint32_t { LONG_MAGIC = 0x533f'f89f, STR_MAGIC = 0xc427'feef, DICT_MAGIC = 0x1ea8'a774 };
 
-enum CertInfoItem : uint32_t {
-  LONG_MAGIC = 0x533f'f89f,
-  STR_MAGIC = 0xc427'feef,
-  DICT_MAGIC = 0x1ea8'a774
-};
+constexpr uint32_t K2_JOB_WORKER_RESPONSE_MAGIC = 0x3afb'3a08;
+constexpr uint32_t CONFDATA_VALUE_MAGIC = 0x3eaa'910b;
 
 } // namespace
 
 namespace tl {
 
 bool K2JobWorkerResponse::fetch(TLBuffer &tlb) noexcept {
-  if (tlb.size() < K2_JOB_WORKER_RESPONSE_MIN_SIZE || *tlb.fetch_trivial<uint32_t>() != K2_JOB_WORKER_RESPONSE_MAGIC) {
+  if (tlb.fetch_trivial<uint32_t>().value_or(TL_ZERO) != K2_JOB_WORKER_RESPONSE_MAGIC) {
     return false;
   }
 
   std::ignore = tlb.fetch_trivial<uint32_t>(); // ignore flags
-  job_id = *tlb.fetch_trivial<int64_t>();
-  const std::string_view body_view{tlb.fetch_string()};
+  const auto opt_job_id{tlb.fetch_trivial<int64_t>()};
+  if (!opt_job_id.has_value()) {
+    return false;
+  }
+  const auto body_view{tlb.fetch_string()};
+
+  job_id = *opt_job_id;
   body = string{body_view.data(), static_cast<string::size_type>(body_view.size())};
   return true;
 }
@@ -130,4 +132,21 @@ bool GetPemCertInfoResponse::fetch(TLBuffer &tlb) noexcept {
   return true;
 }
 
+bool ConfdataValue::fetch(TLBuffer &tlb) noexcept {
+  if (tlb.fetch_trivial<uint32_t>().value_or(TL_ZERO) != CONFDATA_VALUE_MAGIC) {
+    return false;
+  }
+  const auto value_view{tlb.fetch_string()};
+  Bool is_php_serialized_{};
+  Bool is_json_serialized_{};
+  if (!(is_php_serialized_.fetch(tlb) && is_json_serialized_.fetch(tlb))) {
+    return false;
+  }
+
+  value = {value_view.data(), static_cast<string::size_type>(value_view.size())};
+  is_php_serialized = is_php_serialized_.value;
+  is_json_serialized = is_json_serialized_.value;
+  return true;
+}
+
 } // namespace tl