Introduce external storage to store the state of session and operation in Kyuubi Server.

[wForget]

At present, kyuubi server saves some states of session and operation, such as: engine instance, remote handle, etc., in memory, resulting in kyuubi server not completely highly available, and grayscale republish kyuubi server takes a lot of time, so I want to introduce external storage to store the state of session and operation in Kyuubi Server.

---

一、背景

1、需求背景

我们目前的部署模式是在每个集群部署 5 个 Kyuubi Server 服务，使用 Hive JDBC Zookeeper 方式的连接。

此方式存在一些问题：

• 连接不是高可用：服务宕机后，上面已经创建的连接不可用，执行的任务状态丢失，可能会导致重复计算，客户端也需要重新创建连接。
• 服务上线周期长：为了不对线上任务造成影响，我们一般对 Kyuubi Server 使用灰度下线，先在 Zookeeper 上面注销 KyuubiServer，不让有新的连接过来，然后等待 KyuubiServer 上面的连接全部执行完成后才停止 Server服务，再重新上线。对与有较长的离线任务时，整个上线周期特别长。

2、Kyuubi Server 中 SQL 执行流程

2.1、Kyuubi Server 中相关类梳理（伪代码）：

2.2、SQL 执行流程：

• 创建 Connection：使用 Driver（Hive/Kyuubi JDBC Driver）创建一个 Connection ，通过 Zookeeper 获取一个可用的 Kyuubi Server 地址，并创建 Thrift Client。
• 客户端发送 OpenSession 请求：客户端 Connection 初始化时，调用 OpenSession 方法，发生 OpenSession 请求。
• 服务端 OpenSession：Kyuubi Server 接受到 OpenSession 请求后，通过 SessionManager 创建一个 Session 实例，缓存在 handleToSession 的 Map 中，调用 Session 的 open 方法异步（或同步）启动引擎并创建引擎的 Thrift Client，最后向客户端返回 SessionHandle。
• 客户端创建 Statement 并执行：客户端调用 Connection#createStatement() 方法创建一个 Statement，通过 Statement.executeQuery(sql) 执行 SQL，使用 Client 向 Kyuubi Server 发送 ExecuteStatement（带有 SessionHandle） 请求。
• 服务端创建 ExecuteStatement Operation：服务端接受到 ExecuteStatement 请求后，通过 SessionManager 获取 Session，调用 Session 的 executeStatement 方法，通过 OperationManager 创建一个 ExecuteStatement Operation，缓存在 handleToOperation 的 Map 中，执行完 Operation 后，向客户端返回 OperationHandle。
• 服务端执行 ExecuteStatement Operation：Session 的 executeStatement 方法创建好 Operation 后，通过 runOperation 方法执行 Operation，ExecuteStatement Operation 使用 Session 创建的 Client 向 Engine 发送 ExecuteStatement 请求并持有 Remote OperationHandle，启动新的线程监听 Operation 状态。
• 客户端轮询 Operation 状态：客户端使用 OperationHandle 轮询的向服务端发送 GetOperationStatus 请求，获取 Operation 状态，直到返回完成的状态（此过程可获取日志和进度）。
• 客户端 Fetch ResultSet：当有结果集时，ResultSet#next() 使用 client 带上 OperationHandle 发送 FetchResults 请求获取结果集。
• 服务端 FetchResults：服务端根据请求的 OperationHandle 在 OperationManager 中找到对应的 Operation 获取 ResultSet。

3、Kyuubi Server 状态梳理

• Session 对应 Engine：Session 创建后持有 Engine 的 Client，包括了 Engine 的 host、port、remoteSessionHandle
• Operation 对应 Session 和 RemoteOperation：Operation 中关联了 Session 实例，客户端使用 OperationHandle 进行请求，服务端需要根据 OperationHandle 关联 Session 和 Engine 的 Operation，需要：OperationHandle、SessionHandle、RemoteOperationHandle

二、方案

1、实现思路

1.1、服务端

• KVStore：抽象 KV 存储，用于外部存储保存状态，先实现 ZookeeperKVStore。
• SessionExternalStore：外部存储 Session 状态，Key：SessionHandle；Value：EngineInstance(host, port)、RemoteSessionHandle
• OperationExternalStore：外部存储 Operation 状态，Key：OperationHandle；Value：SessionHandle、RemoteOperationHandle、OperationType
• 使用 SessionExternalStore：open session 后通过 SessionExternalStore 保存 session 状态到外部存储中，KyuubiSessionManager 获取 Session 在 handleToSession 中找不到时，再通过 SessionExternalStore 获取 Session
• 使用 OperationExternalStore：仅处理 ExecuteStatement Operation， 在执行 executeStatement 后通过 OperationExternalStore 保存 Operation 状态，KyuubiOperationManager 获取 Operation 在 handleToOperation 中找不到时，再通过 OperationExternalStore 获取 Operation

1.2、客户端

• 可重连的客户端：为 Client 添加 Proxy 类，发生 Thrift 异常时，重新获取 Server 连接建立 Client。

2、实现类抽象

KVStore：Key/Value 存储

trait KVStore {
  def get[T](key: String): Option[T]
  def set(key: String, value: Any): Unit
  def remove(key: String): Unit
}
 
 
class ZookeeperKVStore extends KVStore {
  ...
}

StoreEntity：存储的 Value Entity

trait StoreEntity
 
case class OperationStoreEntity(
    operationHandle: String,
    sessionHandle: String,
    remoteOperationHandle: String) extends StoreEntity
 
case class SessionStoreEntity(
    sessionHandle: String,
    engineSpace: String,
    engineSessionHandle: String) extends StoreEntity

ExternalStore：Session/Operation 外部存储。enity() 方法：转换成存储的 Entity 对象；restore 方法：从 Entity 对象重新生成 Session/Operation 对象。

abstract class ExternalStore[T, E <: StoreEntity](store: KVStore) {
 
  def get(key: String): Option[T] = {
    store.get[E](key).map(restore(_))
  }
 
  def set(key: String, value: T): Unit = {
    store.set(key, entity(value))
  }
 
  def remove(key: String): Unit = store.remove(key)
 
  protected def entity(value: T): E
 
  protected def restore(e: E): T
 
}
 
 
class OperationExternalStore(store: KVStore, sessionManager: KyuubiSessionManager)
  extends ExternalStore[KyuubiOperation, OperationStoreEntity](store) {
  ...
}
 
 
class SessionExternalStore(store: KVStore, sessionManager: KyuubiSessionManager)
  extends ExternalStore[KyuubiSession, SessionStoreEntity](store) {
  ...
}

RestroeSession/Operation：定义 restore 的 Session/Operation 对象，禁用一些不可调用的方法，并进行重新打开连接等操作。

class RestoredKyuubiSessionImpl(protocol: TProtocolVersion,
                                 user: String,
                                 password: String,
                                 ipAddress: String,
                                 conf: Map[String, String],
                                 sessionManager: KyuubiSessionManager,
                                 sessionConf: KyuubiConf)
  extends KyuubiSession(protocol, user, password, ipAddress, conf, sessionManager) {
 
  override val sessionType: SessionType = SessionType.SQL
 
  override def open(): Unit = {
    // TODO throw exception
  }
 
  ...
}
 
 
class RestoredKyuubiOperation(session: Session) extends KyuubiOperation(session) {
 
  override protected def runInternal(): Unit = {
    // TODO throw exception
  }
 
  ...
}

[wForget]

Related issues: #1754

[yaooqinn]

How about extending MetadataManager？