go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/server/quota/internal/lua/utils.lua (about)

     1  -- Copyright 2022 The LUCI Authors
     2  --
     3  -- Licensed under the Apache License, Version 2.0 (the "License");
     4  -- you may not use this file except in compliance with the License.
     5  -- You may obtain a copy of the License at
     6  --
     7  --      http://www.apache.org/licenses/LICENSE-2.0
     8  --
     9  -- Unless required by applicable law or agreed to in writing, software
    10  -- distributed under the License is distributed on an "AS IS" BASIS,
    11  -- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  -- See the License for the specific language governing permissions and
    13  -- limitations under the License.
    14  
    15  -- expect these to be passed in
    16  local PB = ...
    17  assert(PB)
    18  
    19  local redis = redis
    20  
    21  local Utils = {}
    22  
    23  -- Get the current time; We treat the script execution as atomic, so all
    24  -- accounts will be updated at exactly the same timestamp.
    25  --
    26  -- Note that time returns seconds and MICRO seconds (not milliseconds) within
    27  -- that second.
    28  local nowRaw = redis.call('TIME')
    29  Utils.NOW = PB.new('google.protobuf.Timestamp', {
    30    seconds = tonumber(nowRaw[1]),
    31    nanos = tonumber(nowRaw[2]) * 1000,
    32  })
    33  
    34  local math_floor = math.floor
    35  
    36  function Utils.Millis(duration)
    37    if duration == nil then
    38      return nil
    39    end
    40    return math_floor(duration.seconds * 1000 + duration.nanos / 1000000)
    41  end
    42  
    43  function Utils.WithRequestID(req, callback)
    44    local req_id_key = req.request_key
    45  
    46    if req_id_key ~= "" then
    47      local hash_scheme, hash, response = unpack(redis.call(
    48         "HMGET", req_id_key, "hash_scheme", "hash", "response"))
    49      if hash_scheme ~= nil then
    50        local my_scheme = tonumber(hash_scheme)
    51        if response then
    52          if req.hash_scheme ~= my_scheme then
    53            -- hash scheme mismatch; just return the response
    54            return response
    55          elseif req.hash_scheme == my_scheme and req.hash == hash then
    56            -- hash match; return response.
    57            return response
    58          else
    59            return redis.error_reply('REQUEST_HASH')
    60          end
    61        end
    62      end
    63    end
    64  
    65    -- either there was no request_key, or this is the first call.
    66    local ret, allOK = callback()
    67  
    68    local response = PB.marshal(ret)
    69  
    70    if allOK and req_id_key ~= "" then
    71      redis.call(
    72        "HSET", req_id_key,
    73        "hash_scheme", tostring(req.hash_scheme),
    74        "hash", req.hash,
    75        "response", response
    76      )
    77      local ttl = Utils.Millis(req.request_key_ttl)
    78      if ttl == nil then
    79        -- 2 hours * 60 minutes * 60 seconds * 1000 milliseconds.
    80        ttl = 2 * 60 * 60 * 1000
    81      end
    82      if ttl > 0 then
    83        redis.call("PEXPIRE", req_id_key, ttl)
    84      end
    85    end
    86  
    87    return response
    88  end
    89  
    90  return Utils