diff --git a/src/main/java/org/redkale/net/client/Client.java b/src/main/java/org/redkale/net/client/Client.java index 5b435443a..d285749bd 100644 --- a/src/main/java/org/redkale/net/client/Client.java +++ b/src/main/java/org/redkale/net/client/Client.java @@ -81,7 +81,7 @@ public abstract class Client, R extends ClientR protected Supplier closeRequestSupplier; // 创建连接后进行的登录鉴权操作 - protected Function>> authenticate; + protected BiFunction>> authenticate; protected Client(String name, AsyncGroup group, ClientAddress address) { this(name, group, true, address, Utility.cpus(), DEFAULT_MAX_PIPELINES, null, null, null); @@ -106,7 +106,7 @@ public abstract class Client, R extends ClientR boolean tcp, ClientAddress address, int maxConns, - Function>> authenticate) { + BiFunction>> authenticate) { this(name, group, tcp, address, maxConns, DEFAULT_MAX_PIPELINES, null, null, authenticate); } @@ -117,7 +117,7 @@ public abstract class Client, R extends ClientR ClientAddress address, int maxConns, Supplier closeRequestSupplier, - Function>> authenticate) { + BiFunction>> authenticate) { this(name, group, tcp, address, maxConns, DEFAULT_MAX_PIPELINES, null, closeRequestSupplier, authenticate); } @@ -131,7 +131,7 @@ public abstract class Client, R extends ClientR int maxPipelines, Supplier pingRequestSupplier, Supplier closeRequestSupplier, - Function>> authenticate) { + BiFunction>> authenticate) { if (maxPipelines < 1) { throw new IllegalArgumentException("maxPipelines must bigger 0"); } @@ -346,6 +346,9 @@ public abstract class Client, R extends ClientR R virtualReq = createVirtualRequestAfterConnect(); if (virtualReq != null) { virtualReq.traceid = traceid; + if (virtualReq.workThread == null) { + virtualReq.workThread = workThread; + } future = future.thenCompose(conn -> { Traces.currentTraceid(traceid); return conn.writeVirtualRequest(virtualReq).thenApply(v -> conn); @@ -357,7 +360,7 @@ public abstract class Client, R extends ClientR }); } if (authenticate != null) { - future = future.thenCompose(authenticate.apply(traceid)); + future = future.thenCompose(authenticate.apply(workThread, traceid)); } return future.thenApply(c -> { c.setAuthenticated(true);