parquet/encryption/two_level_cache_with_expiration.h

0001 // Licensed to the Apache Software Foundation (ASF) under one
0002 // or more contributor license agreements.  See the NOTICE file
0003 // distributed with this work for additional information
0004 // regarding copyright ownership.  The ASF licenses this file
0005 // to you under the Apache License, Version 2.0 (the
0006 // "License"); you may not use this file except in compliance
0007 // with the License.  You may obtain a copy of the License at
0008 //
0009 //   http://www.apache.org/licenses/LICENSE-2.0
0010 //
0011 // Unless required by applicable law or agreed to in writing,
0012 // software distributed under the License is distributed on an
0013 // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
0014 // KIND, either express or implied.  See the License for the
0015 // specific language governing permissions and limitations
0016 // under the License.
0017
0018 #pragma once
0019
0020 #include <chrono>
0021 #include <unordered_map>
0022
0023 #include "arrow/util/concurrent_map.h"
0024 #include "arrow/util/mutex.h"
0025
0026 namespace parquet::encryption {
0027
0028 using ::arrow::util::ConcurrentMap;
0029
0030 namespace internal {
0031
0032 using TimePoint =
0033     std::chrono::time_point<std::chrono::system_clock, std::chrono::duration<double>>;
0034
0035 inline TimePoint CurrentTimePoint() { return std::chrono::system_clock::now(); }
0036
0037 template <typename E>
0038 class ExpiringCacheEntry {
0039  public:
0040   ExpiringCacheEntry() = default;
0041
0042   ExpiringCacheEntry(E cached_item, double expiration_interval_seconds)
0043       : expiration_timestamp_(CurrentTimePoint() +
0044                               std::chrono::duration<double>(expiration_interval_seconds)),
0045         cached_item_(std::move(cached_item)) {}
0046
0047   bool IsExpired() const {
0048     const auto now = CurrentTimePoint();
0049     return (now > expiration_timestamp_);
0050   }
0051
0052   E cached_item() { return cached_item_; }
0053
0054  private:
0055   const TimePoint expiration_timestamp_;
0056   E cached_item_;
0057 };
0058
0059 // This class is to avoid the below warning when compiling KeyToolkit class with VS2015
0060 // warning C4503: decorated name length exceeded, name was truncated
0061 template <typename V>
0062 class ExpiringCacheMapEntry {
0063  public:
0064   ExpiringCacheMapEntry() = default;
0065
0066   explicit ExpiringCacheMapEntry(
0067       std::shared_ptr<ConcurrentMap<std::string, V>> cached_item,
0068       double expiration_interval_seconds)
0069       : map_cache_(cached_item, expiration_interval_seconds) {}
0070
0071   bool IsExpired() { return map_cache_.IsExpired(); }
0072
0073   std::shared_ptr<ConcurrentMap<std::string, V>> cached_item() {
0074     return map_cache_.cached_item();
0075   }
0076
0077  private:
0078   // ConcurrentMap object may be accessed and modified at many places at the same time,
0079   // from multiple threads, or even removed from cache.
0080   ExpiringCacheEntry<std::shared_ptr<ConcurrentMap<std::string, V>>> map_cache_;
0081 };
0082
0083 }  // namespace internal
0084
0085 // Two-level cache with expiration of internal caches according to token lifetime.
0086 // External cache is per token, internal is per string key.
0087 // Wrapper class around:
0088 //    std::unordered_map<std::string,
0089 //    internal::ExpiringCacheEntry<std::unordered_map<std::string, V>>>
0090 // This cache is safe to be shared between threads.
0091 template <typename V>
0092 class TwoLevelCacheWithExpiration {
0093  public:
0094   TwoLevelCacheWithExpiration() {
0095     last_cache_cleanup_timestamp_ = internal::CurrentTimePoint();
0096   }
0097
0098   std::shared_ptr<ConcurrentMap<std::string, V>> GetOrCreateInternalCache(
0099       const std::string& access_token, double cache_entry_lifetime_seconds) {
0100     auto lock = mutex_.Lock();
0101
0102     auto external_cache_entry = cache_.find(access_token);
0103     if (external_cache_entry == cache_.end() ||
0104         external_cache_entry->second.IsExpired()) {
0105       cache_.insert({access_token, internal::ExpiringCacheMapEntry<V>(
0106                                        std::make_shared<ConcurrentMap<std::string, V>>(),
0107                                        cache_entry_lifetime_seconds)});
0108     }
0109
0110     return cache_[access_token].cached_item();
0111   }
0112
0113   void CheckCacheForExpiredTokens(double cache_cleanup_period_seconds = 0.0) {
0114     auto lock = mutex_.Lock();
0115
0116     const auto now = internal::CurrentTimePoint();
0117     if (now > (last_cache_cleanup_timestamp_ +
0118                std::chrono::duration<double>(cache_cleanup_period_seconds))) {
0119       RemoveExpiredEntriesNoMutex();
0120       last_cache_cleanup_timestamp_ = now;
0121     }
0122   }
0123
0124   void Remove(const std::string& access_token) {
0125     auto lock = mutex_.Lock();
0126     cache_.erase(access_token);
0127   }
0128
0129   void Clear() {
0130     auto lock = mutex_.Lock();
0131     cache_.clear();
0132   }
0133
0134  private:
0135   void RemoveExpiredEntriesNoMutex() {
0136     for (auto it = cache_.begin(); it != cache_.end();) {
0137       if (it->second.IsExpired()) {
0138         it = cache_.erase(it);
0139       } else {
0140         ++it;
0141       }
0142     }
0143   }
0144   std::unordered_map<std::string, internal::ExpiringCacheMapEntry<V>> cache_;
0145   internal::TimePoint last_cache_cleanup_timestamp_;
0146   ::arrow::util::Mutex mutex_;
0147 };
0148
0149 }  // namespace parquet::encryption