Implement FUNCTION commands like redis

src/commands/cmd_function.cc

@@ -0,0 +1,100 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ *
+ */
+
+#include "commander.h"
+#include "commands/command_parser.h"
+#include "parse_util.h"
+#include "server/redis_reply.h"
+#include "storage/scripting.h"
+#include "string_util.h"
+
+namespace redis {
+
+struct CommandFunction : Commander {
+  Status Execute(Server *srv, Connection *conn, std::string *output) override {
+    CommandParser parser(args_, 1);
+    if (parser.EatEqICase("load")) {
+      bool replace = false;
+      if (parser.EatEqICase("replace")) {
+        replace = true;
+      }
+
+      std::string libname;
+      auto s = lua::FunctionLoad(srv, GET_OR_RET(parser.TakeStr()), true, replace, &libname);
+      if (!s) return s;
+
+      *output = SimpleString(libname);
+      return Status::OK();
+    } else if (parser.EatEqICase("list")) {
+      std::string libname;
+      if (parser.EatEqICase("libraryname")) {
+        libname = GET_OR_RET(parser.TakeStr());
+      }
+
+      bool with_code = false;
+      if (parser.EatEqICase("withcode")) {
+        with_code = true;
+      }
+
+      return lua::FunctionList(srv, libname, with_code, output);
+    } else if (parser.EatEqICase("listfunc")) {
+      std::string funcname;
+      if (parser.EatEqICase("funcname")) {
+        funcname = GET_OR_RET(parser.TakeStr());
+      }
+
+      return lua::FunctionListFunc(srv, funcname, output);
+    } else if (parser.EatEqICase("delete")) {
+      auto libname = GET_OR_RET(parser.TakeStr());
+      if (!lua::FunctionIsLibExist(srv, libname)) {
+        return {Status::NotOK, "no such library"};
+      }
+
+      auto s = lua::FunctionDelete(srv, libname);
+      if (!s) return s;
+
+      *output = SimpleString("OK");
+      return Status::OK();
+    } else {
+      return {Status::NotOK, "no such subcommand"};
+    }
+  }
+};
+
+struct CommandFCall : Commander {
+  Status Execute(Server *srv, Connection *conn, std::string *output) override {
+    int64_t numkeys = GET_OR_RET(ParseInt<int64_t>(args_[2], 10));
+    if (numkeys > int64_t(args_.size() - 3)) {
+      return {Status::NotOK, "Number of keys can't be greater than number of args"};
+    } else if (numkeys < -1) {
+      return {Status::NotOK, "Number of keys can't be negative"};
+    }
+
+    return lua::FunctionCall(srv, args_[1], std::vector<std::string>(args_.begin() + 3, args_.begin() + 3 + numkeys),
+                             std::vector<std::string>(args_.begin() + 3 + numkeys, args_.end()), output);
+  }
+};
+
+CommandKeyRange GetScriptEvalKeyRange(const std::vector<std::string> &args);
+
+REDIS_REGISTER_COMMANDS(MakeCmdAttr<CommandFunction>("function", -2, "exclusive no-script", 0, 0, 0),
+                        MakeCmdAttr<CommandFCall>("fcall", -3, "exclusive write no-script", GetScriptEvalKeyRange));
+
+}  // namespace redis

src/server/server.cc

@@ -44,6 +44,7 @@
 #include "storage/compaction_checker.h"
 #include "storage/redis_db.h"
 #include "storage/scripting.h"
+#include "storage/storage.h"
 #include "string_util.h"
 #include "thread_util.h"
 #include "time_util.h"
@@ -1527,7 +1528,7 @@ Status Server::ScriptExists(const std::string &sha) {
 }
 
 Status Server::ScriptGet(const std::string &sha, std::string *body) const {
-  std::string func_name = engine::kLuaFunctionPrefix + sha;
+  std::string func_name = engine::kLuaFuncSHAPrefix + sha;
   auto cf = storage->GetCFHandle(engine::kPropagateColumnFamilyName);
   auto s = storage->Get(rocksdb::ReadOptions(), cf, func_name, body);
   if (!s.ok()) {
@@ -1537,10 +1538,40 @@ Status Server::ScriptGet(const std::string &sha, std::string *body) const {
 }
 
 Status Server::ScriptSet(const std::string &sha, const std::string &body) const {
-  std::string func_name = engine::kLuaFunctionPrefix + sha;
+  std::string func_name = engine::kLuaFuncSHAPrefix + sha;
   return storage->WriteToPropagateCF(func_name, body);
 }
 
+Status Server::FunctionGetCode(const std::string &lib, std::string *code) const {
+  std::string func_name = engine::kLuaLibCodePrefix + lib;
+  auto cf = storage->GetCFHandle(engine::kPropagateColumnFamilyName);
+  auto s = storage->Get(rocksdb::ReadOptions(), cf, func_name, code);
+  if (!s.ok()) {
+    return {s.IsNotFound() ? Status::NotFound : Status::NotOK, s.ToString()};
+  }
+  return Status::OK();
+}
+
+Status Server::FunctionGetLib(const std::string &func, std::string *lib) const {
+  std::string func_name = engine::kLuaFuncLibPrefix + func;
+  auto cf = storage->GetCFHandle(engine::kPropagateColumnFamilyName);
+  auto s = storage->Get(rocksdb::ReadOptions(), cf, func_name, lib);
+  if (!s.ok()) {
+    return {s.IsNotFound() ? Status::NotFound : Status::NotOK, s.ToString()};
+  }
+  return Status::OK();
+}
+
+Status Server::FunctionSetCode(const std::string &lib, const std::string &code) const {
+  std::string func_name = engine::kLuaLibCodePrefix + lib;
+  return storage->WriteToPropagateCF(func_name, code);
+}
+
+Status Server::FunctionSetLib(const std::string &func, const std::string &lib) const {
+  std::string func_name = engine::kLuaFuncLibPrefix + func;
+  return storage->WriteToPropagateCF(func_name, lib);
+}
+
 void Server::ScriptReset() {
   auto lua = lua_.exchange(lua::CreateState(this));
   lua::DestroyState(lua);

src/server/server.h

@@ -255,6 +255,11 @@ class Server {
   void ScriptReset();
   Status ScriptFlush();
 
+  Status FunctionGetCode(const std::string &lib, std::string *code) const;
+  Status FunctionGetLib(const std::string &func, std::string *lib) const;
+  Status FunctionSetCode(const std::string &lib, const std::string &code) const;
+  Status FunctionSetLib(const std::string &func, const std::string &lib) const;
+
   Status Propagate(const std::string &channel, const std::vector<std::string> &tokens) const;
   Status ExecPropagatedCommand(const std::vector<std::string> &tokens);
   Status ExecPropagateScriptCommand(const std::vector<std::string> &tokens);

src/storage/redis_db.h

@@ -66,20 +66,20 @@ class Database {
   std::string namespace_;
 
   friend class LatestSnapShot;
-  class LatestSnapShot {
-   public:
-    explicit LatestSnapShot(engine::Storage *storage)
-        : storage_(storage), snapshot_(storage_->GetDB()->GetSnapshot()) {}
-    ~LatestSnapShot() { storage_->GetDB()->ReleaseSnapshot(snapshot_); }
-    const rocksdb::Snapshot *GetSnapShot() const { return snapshot_; }
+};
+
+class LatestSnapShot {
+ public:
+  explicit LatestSnapShot(engine::Storage *storage) : storage_(storage), snapshot_(storage_->GetDB()->GetSnapshot()) {}
+  ~LatestSnapShot() { storage_->GetDB()->ReleaseSnapshot(snapshot_); }
+  const rocksdb::Snapshot *GetSnapShot() const { return snapshot_; }
 
-    LatestSnapShot(const LatestSnapShot &) = delete;
-    LatestSnapShot &operator=(const LatestSnapShot &) = delete;
+  LatestSnapShot(const LatestSnapShot &) = delete;
+  LatestSnapShot &operator=(const LatestSnapShot &) = delete;
 
-   private:
-    engine::Storage *storage_ = nullptr;
-    const rocksdb::Snapshot *snapshot_ = nullptr;
-  };
+ private:
+  engine::Storage *storage_ = nullptr;
+  const rocksdb::Snapshot *snapshot_ = nullptr;
 };
 
 class SubKeyScanner : public redis::Database {