From 814da9068c157f9fe9194ba2e38fc3aa2fed51e3 Mon Sep 17 00:00:00 2001
From: wangjiawei04 <wangjiawei04@baidu.com>
Date: Tue, 25 Jun 2019 06:52:28 +0800
Subject: [PATCH] add kvdb to resource

Change-Id: I73890ecaffe1d500dfc66c76d83b0fa8e64732e5
---
 demo-serving/op/kvdb_echo_op.cpp | 10 +++-------
 demo-serving/op/kvdb_echo_op.h   |  5 +----
 predictor/CMakeLists.txt         |  6 ++++--
 predictor/framework/resource.cpp | 13 ++++++++++++-
 predictor/framework/resource.h   |  7 ++++++-
 5 files changed, 26 insertions(+), 15 deletions(-)

diff --git a/demo-serving/op/kvdb_echo_op.cpp b/demo-serving/op/kvdb_echo_op.cpp
index b847acde..1c4d013a 100644
--- a/demo-serving/op/kvdb_echo_op.cpp
+++ b/demo-serving/op/kvdb_echo_op.cpp
@@ -23,13 +23,13 @@ using baidu::paddle_serving::predictor::format::KVDBRes;
 using baidu::paddle_serving::predictor::echo_kvdb_service::Request;
 using baidu::paddle_serving::predictor::echo_kvdb_service::Response;
 
-std::shared_ptr<RocksDBWrapper> KVDBEchoOp::db;
 int KVDBEchoOp::inference() {
     debug();
 }
 int KVDBEchoOp::debug() {
    //TODO: implement DEBUG mode
-    this->DBInit();
+    baidu::paddle_serving::predictor::Resource& resource = baidu::paddle_serving::predictor::Resource::instance();
+    std::shared_ptr<RocksDBWrapper> db = resource.getDB();
     const Request* req = dynamic_cast<const Request*>(get_request_message());
     Response* res = mutable_data<Response>();
     LOG(INFO) << "Receive request in KVDB echo service: " << req->ShortDebugString();
@@ -51,11 +51,7 @@ int KVDBEchoOp::debug() {
     return 0;
 }
 
-void KVDBEchoOp::DBInit() {
-    if (db.get() == nullptr) {
-        db = RocksDBWrapper::RocksDBWrapperFactory("kvdb");
-    }
-}
+
 
 DEFINE_OP(KVDBEchoOp);
 }
diff --git a/demo-serving/op/kvdb_echo_op.h b/demo-serving/op/kvdb_echo_op.h
index de8c5170..7c7502a1 100644
--- a/demo-serving/op/kvdb_echo_op.h
+++ b/demo-serving/op/kvdb_echo_op.h
@@ -14,7 +14,7 @@
 
 #pragma once
 #include "demo-serving/echo_kvdb_service.pb.h"
-
+#include "predictor/framework/resource.h"
 #include "predictor/common/inner_common.h"
 #include "predictor/framework/channel.h"
 #include "predictor/framework/op_repository.h"
@@ -30,9 +30,6 @@ class KVDBEchoOp: public OpWithChannel<baidu::paddle_serving::predictor::echo_kv
       DECLARE_OP(KVDBEchoOp);
       int inference();
       int debug();
-  protected:
-      static std::shared_ptr<RocksDBWrapper> db;
-      void DBInit();
       
 };
 }  //  namespace predictor
diff --git a/predictor/CMakeLists.txt b/predictor/CMakeLists.txt
index c70c5dbe..65221144 100644
--- a/predictor/CMakeLists.txt
+++ b/predictor/CMakeLists.txt
@@ -6,15 +6,17 @@ include(framework/CMakeLists.txt)
 #include(plugin/CMakeLists.txt)
 include(src/CMakeLists.txt)
 
+include_directories(SYSTEM  ${CMAKE_CURRENT_LIST_DIR}/../kvdb/include)
+
 add_library(pdserving ${pdserving_srcs})
 set_source_files_properties(
         ${pdserving_srcs}
         PROPERTIES
         COMPILE_FLAGS  "-Wno-strict-aliasing -Wno-unused-variable -Wno-non-virtual-dtor -Wno-error=non-virtual-dtor -Wno-error=delete-non-virtual-dtor")
-add_dependencies(pdserving protobuf boost brpc leveldb pdcodegen configure)
+add_dependencies(pdserving protobuf kvdb boost brpc leveldb pdcodegen configure)
 
 target_link_libraries(pdserving
-        brpc protobuf boost leveldb configure -lpthread -lcrypto -lm -lrt -lssl -ldl -lz)
+        brpc protobuf boost leveldb configure kvdb -lpthread -lcrypto -lm -lrt -lssl -ldl -lz)
 
 # install
 install(TARGETS pdserving
diff --git a/predictor/framework/resource.cpp b/predictor/framework/resource.cpp
index 9ced89ff..9b545c91 100644
--- a/predictor/framework/resource.cpp
+++ b/predictor/framework/resource.cpp
@@ -36,7 +36,13 @@ DynamicResource::DynamicResource() {}
 
 DynamicResource::~DynamicResource() {}
 
-int DynamicResource::initialize() { return 0; }
+int DynamicResource::initialize() { 
+    return 0;
+}
+
+std::shared_ptr<RocksDBWrapper> Resource::getDB() {
+    return db;
+}
 
 int DynamicResource::clear() { return 0; }
 
@@ -80,6 +86,11 @@ int Resource::initialize(const std::string& path, const std::string& file) {
     LOG(ERROR) << "unable to create tls_bthread_key of thrd_data";
     return -1;
   }
+    //init rocksDB instance
+  if (db.get() == nullptr) {
+    db = RocksDBWrapper::RocksDBWrapperFactory("kvdb");
+  }
+
   THREAD_SETSPECIFIC(_tls_bspec_key, NULL);
   return 0;
 }
diff --git a/predictor/framework/resource.h b/predictor/framework/resource.h
index 3f3bdb91..767b71d4 100644
--- a/predictor/framework/resource.h
+++ b/predictor/framework/resource.h
@@ -16,6 +16,7 @@
 #include <string>
 #include "predictor/common/inner_common.h"
 #include "predictor/framework/memory.h"
+#include "kvdb/paddle_rocksdb.h"
 
 namespace baidu {
 namespace paddle_serving {
@@ -30,6 +31,7 @@ struct DynamicResource {
   int initialize();
 
   int clear();
+
 };
 
 class Resource {
@@ -53,6 +55,8 @@ class Resource {
 
   int finalize();
 
+  std::shared_ptr<RocksDBWrapper> getDB();
+
   DynamicResource* get_dynamic_resource() {
     return reinterpret_cast<DynamicResource*>(
         THREAD_GETSPECIFIC(_tls_bspec_key));
@@ -60,7 +64,8 @@ class Resource {
 
  private:
   int thread_finalize() { return 0; }
-
+  std::shared_ptr<RocksDBWrapper> db;
+  
   THREAD_KEY_T _tls_bspec_key;
 };
 
-- 
GitLab