From 2a903b01fe4dcd4756b5cec86669d0eb85158f61 Mon Sep 17 00:00:00 2001 From: CHzxp <1959044956@qq.com> Date: Tue, 6 Dec 2022 16:07:40 +0800 Subject: [PATCH] [KYUUBI #3910] Add delta parameter for getAndIncrement method ### _Why are the changes needed?_ getAndIncrement method only increase index before this pr, add a delta parameter can increase and decrease index. this pr is mainly for pr KYUUBI #3695 ### _How was this patch tested?_ - [ ] Add some test cases that check the changes thoroughly including negative and positive cases if possible - [ ] Add screenshots for manual tests if appropriate - [ ] [Run test](https://kyuubi.apache.org/docs/latest/develop_tools/testing.html#running-tests) locally before make a pull request Closes #3910 from CHzxp/getAndIncrement_branch. Closes #3910 5c0cca85 [CHzxp] fix bug 3514471e [CHzxp] modify imp f1ae57d6 [xinping] add delta parameter for getAndIncrement Lead-authored-by: CHzxp <1959044956@qq.com> Co-authored-by: xinping Signed-off-by: Cheng Pan --- .../scala/org/apache/kyuubi/ha/client/DiscoveryClient.scala | 3 ++- .../apache/kyuubi/ha/client/etcd/EtcdDiscoveryClient.scala | 4 ++-- .../kyuubi/ha/client/zookeeper/ZookeeperDiscoveryClient.scala | 4 ++-- 3 files changed, 6 insertions(+), 5 deletions(-) diff --git a/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/DiscoveryClient.scala b/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/DiscoveryClient.scala index 29c1734f8..77588180e 100644 --- a/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/DiscoveryClient.scala +++ b/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/DiscoveryClient.scala @@ -174,9 +174,10 @@ trait DiscoveryClient extends Logging { * Atomically get an Int number and add one * @param path the path of stored data, * If the path does not exist, it will be created and initialized to 0 + * @param delta the increase num * @return the stored data under path */ - def getAndIncrement(path: String): Int + def getAndIncrement(path: String, delta: Int = 1): Int } object DiscoveryClient { diff --git a/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/etcd/EtcdDiscoveryClient.scala b/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/etcd/EtcdDiscoveryClient.scala index 727ba5cae..dabdc7063 100644 --- a/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/etcd/EtcdDiscoveryClient.scala +++ b/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/etcd/EtcdDiscoveryClient.scala @@ -303,7 +303,7 @@ class EtcdDiscoveryClient(conf: KyuubiConf) extends DiscoveryClient { ByteSequence.from(initData.getBytes())).get() } - def getAndIncrement(path: String): Int = { + def getAndIncrement(path: String, delta: Int = 1): Int = { val lockPath = s"${path}_tmp_for_lock" tryWithLock(lockPath, 60 * 1000) { if (pathNonExists(path)) { @@ -311,7 +311,7 @@ class EtcdDiscoveryClient(conf: KyuubiConf) extends DiscoveryClient { setData(path, String.valueOf(0).getBytes) } val s = new String(getData(path)).toInt - setData(path, String.valueOf(s + 1).getBytes) + setData(path, String.valueOf(s + delta).getBytes) s } } diff --git a/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/zookeeper/ZookeeperDiscoveryClient.scala b/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/zookeeper/ZookeeperDiscoveryClient.scala index 4802b4f7b..ca46d1130 100644 --- a/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/zookeeper/ZookeeperDiscoveryClient.scala +++ b/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/zookeeper/ZookeeperDiscoveryClient.scala @@ -304,11 +304,11 @@ class ZookeeperDiscoveryClient(conf: KyuubiConf) extends DiscoveryClient { secretNode.start() } - def getAndIncrement(path: String): Int = { + def getAndIncrement(path: String, delta: Int = 1): Int = { val dai = new DistributedAtomicInteger(zkClient, path, new RetryForever(1000)) var atomicVal: AtomicValue[Integer] = null do { - atomicVal = dai.increment() + atomicVal = dai.add(delta) } while (atomicVal == null || !atomicVal.succeeded()) atomicVal.preValue().intValue() }