Status EvalGenericCommand()

in src/storage/scripting.cc [608:714]

• 74 lines of code
• 11 McCabe index (conditional complexity)

Status EvalGenericCommand(redis::Connection *conn, engine::Context *ctx, const std::string &body_or_sha,
                          const std::vector<std::string> &keys, const std::vector<std::string> &argv, bool evalsha,
                          std::string *output, bool read_only) {
  Server *srv = conn->GetServer();
  // Use the worker's private Lua VM when entering the read-only mode
  lua_State *lua = conn->Owner()->Lua();

  /* We obtain the script SHA1, then check if this function is already
   * defined into the Lua state */
  char funcname[2 + 40 + 1] = {};
  memcpy(funcname, REDIS_LUA_FUNC_SHA_PREFIX, sizeof(REDIS_LUA_FUNC_SHA_PREFIX));

  if (!evalsha) {
    SHA1Hex(funcname + 2, body_or_sha.c_str(), body_or_sha.size());
  } else {
    for (int j = 0; j < 40; j++) {
      funcname[j + 2] = static_cast<char>(tolower(body_or_sha[j]));
    }
  }

  /* Push the pcall error handler function on the stack. */
  lua_getglobal(lua, "__redis__err__handler");

  /* Try to lookup the Lua function */
  lua_getglobal(lua, funcname);
  if (lua_isnil(lua, -1)) {
    lua_pop(lua, 1); /* remove the nil from the stack */
    std::string body;
    if (evalsha) {
      auto s = srv->ScriptGet(funcname + 2, &body);
      if (!s.IsOK()) {
        lua_pop(lua, 1); /* remove the error handler from the stack. */
        return {Status::RedisNoScript, redis::errNoMatchingScript};
      }
    } else {
      body = body_or_sha;
    }

    std::string sha = funcname + 2;
    auto s = CreateFunction(srv, body, &sha, lua, false);
    if (!s.IsOK()) {
      lua_pop(lua, 1); /* remove the error handler from the stack. */
      return s;
    }
    /* Now the following is guaranteed to return non nil */
    lua_getglobal(lua, funcname);
  }

  ScriptRunCtx current_script_run_ctx;
  current_script_run_ctx.conn = conn;
  current_script_run_ctx.ctx = ctx;
  current_script_run_ctx.flags = read_only ? ScriptFlagType::kScriptNoWrites : 0;
  lua_getglobal(lua, fmt::format(REDIS_LUA_FUNC_SHA_FLAGS, funcname + 2).c_str());
  if (!lua_isnil(lua, -1)) {
    // It should be ensured that the conversion is successful
    auto script_flags = lua_tointeger(lua, -1);
    current_script_run_ctx.flags |= script_flags;
  }
  lua_pop(lua, 1);

  SaveOnRegistry(lua, REGISTRY_SCRIPT_RUN_CTX_NAME, &current_script_run_ctx);

  // For the Lua script, should be always run with RESP2 protocol,
  // unless users explicitly set the protocol version in the script via `redis.setresp`.
  // So we need to save the current protocol version and set it to RESP2,
  // and then restore it after the script execution.
  auto saved_protocol_version = conn->GetProtocolVersion();
  conn->SetProtocolVersion(redis::RESP::v2);
  /* Populate the argv and keys table accordingly to the arguments that
   * EVAL received. */
  SetGlobalArray(lua, "KEYS", keys);
  SetGlobalArray(lua, "ARGV", argv);

  if (lua_pcall(lua, 0, 1, -2)) {
    auto msg = fmt::format("running script (call to {}): {}", funcname, lua_tostring(lua, -1));
    *output = redis::Error({Status::NotOK, msg});
    lua_pop(lua, 2);
  } else {
    *output = ReplyToRedisReply(conn, lua);
    lua_pop(lua, 2);
  }
  conn->SetProtocolVersion(saved_protocol_version);

  // clean global variables to prevent information leak in function commands
  lua_pushnil(lua);
  lua_setglobal(lua, "KEYS");
  lua_pushnil(lua);
  lua_setglobal(lua, "ARGV");
  RemoveFromRegistry(lua, REGISTRY_SCRIPT_RUN_CTX_NAME);

  /* Call the Lua garbage collector from time to time to avoid a
   * full cycle performed by Lua, which adds too latency.
   *
   * The call is performed every LUA_GC_CYCLE_PERIOD executed commands
   * (and for LUA_GC_CYCLE_PERIOD collection steps) because calling it
   * for every command uses too much CPU. */
  constexpr int64_t LUA_GC_CYCLE_PERIOD = 50;
  static int64_t gc_count = 0;

  gc_count++;
  if (gc_count == LUA_GC_CYCLE_PERIOD) {
    lua_gc(lua, LUA_GCSTEP, LUA_GC_CYCLE_PERIOD);
    gc_count = 0;
  }

  return Status::OK();
}