optimize NebulaSessionClient layer

nebula-contrib · Oct 10, 2023 · c485509 · c485509
1 parent 4b69559
commit c485509
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 34 deletions.
diff --git a/core/src/main/scala/zio/nebula/NebulaSessionClient.scala b/core/src/main/scala/zio/nebula/NebulaSessionClient.scala
@@ -55,35 +55,30 @@ trait NebulaSessionClient {
 
 object NebulaSessionClient {
 
-  def sessionLayer: ZLayer[NebulaSessionPoolConfig & Scope, Throwable, SessionPool] =
-    ZLayer.fromZIO {
-      ZIO.serviceWithZIO[NebulaSessionPoolConfig](nebulaConfig =>
-        ZIO.acquireRelease(
-          ZIO.attempt(
-            new SessionPool(
-              new SessionPoolConfig(
-                nebulaConfig.address.map(d => new HostAddress(d.host, d.port)).asJava,
-                nebulaConfig.spaceName,
-                nebulaConfig.auth.username,
-                nebulaConfig.auth.password
-              ).setMaxSessionSize(nebulaConfig.maxSessionSize)
-                .setMinSessionSize(nebulaConfig.minSessionSize)
-                .setRetryTimes(nebulaConfig.retryTimes)
-                .setWaitTime(nebulaConfig.waitTimeMills)
-                .setIntervalTime(nebulaConfig.intervalTimeMills)
-                .setTimeout(nebulaConfig.timeoutMills)
-                .setCleanTime(nebulaConfig.cleanTimeSeconds)
-                .setReconnect(nebulaConfig.reconnect)
-                .setHealthCheckTime(nebulaConfig.healthCheckTimeSeconds)
-                .setUseHttp2(nebulaConfig.useHttp2)
-            )
-          )
-        )(release => ZIO.attempt(release.close()).onError(e => ZIO.logErrorCause(e)).ignoreLogged)
-      )
-    }
-
-  lazy val layer: ZLayer[NebulaSessionPoolConfig & Scope, Throwable, NebulaSessionClient] =
-    sessionLayer >>> ZLayer.fromZIO(
-      ZIO.serviceWith[SessionPool](new NebulaSessionClientLive(_))
-    )
+  lazy val layer: ZLayer[Scope with NebulaSessionPoolConfig, Throwable, NebulaSessionClient] = ZLayer.fromZIO {
+    for {
+      nebulaConfig <- ZIO.service[NebulaSessionPoolConfig]
+      sessionPool  <- ZIO.acquireRelease(
+                        ZIO.attempt(
+                          new SessionPool(
+                            new SessionPoolConfig(
+                              nebulaConfig.address.map(d => new HostAddress(d.host, d.port)).asJava,
+                              nebulaConfig.spaceName,
+                              nebulaConfig.auth.username,
+                              nebulaConfig.auth.password
+                            ).setMaxSessionSize(nebulaConfig.maxSessionSize)
+                              .setMinSessionSize(nebulaConfig.minSessionSize)
+                              .setRetryTimes(nebulaConfig.retryTimes)
+                              .setWaitTime(nebulaConfig.waitTimeMills)
+                              .setIntervalTime(nebulaConfig.intervalTimeMills)
+                              .setTimeout(nebulaConfig.timeoutMills)
+                              .setCleanTime(nebulaConfig.cleanTimeSeconds)
+                              .setReconnect(nebulaConfig.reconnect)
+                              .setHealthCheckTime(nebulaConfig.healthCheckTimeSeconds)
+                              .setUseHttp2(nebulaConfig.useHttp2)
+                          )
+                        )
+                      )(release => ZIO.attempt(release.close()).ignoreLogged)
+    } yield new NebulaSessionClientLive(sessionPool)
+  }
 }
diff --git a/core/src/test/scala/zio/nebula/ZioNebulaEnvironment.scala b/core/src/test/scala/zio/nebula/ZioNebulaEnvironment.scala
@@ -1,13 +1,10 @@
 package zio.nebula
 
 import zio._
-import zio.nebula.NebulaSessionClient.sessionLayer
 import zio.nebula.meta.NebulaMetaClient
 import zio.nebula.net.NebulaClient
 import zio.nebula.storage.NebulaStorageClient
 
-import com.vesoft.nebula.client.graph.SessionPool
-
 /**
  * This is the default configuration dedicated to testing.
  *