[dingo-client,dingo-license] Fixing license check and client vector upsert

wt0530 · ketor · commit 7f5d4cc252ba · 2024-05-31T18:36:22.000+08:00
diff --git a/dingo-client/src/main/java/io/dingodb/client/OperationServiceV2.java b/dingo-client/src/main/java/io/dingodb/client/OperationServiceV2.java
@@ -119,6 +119,7 @@ public OperationServiceV2(String coordinatorSvr) {
 
     public void close() {
         jobManager.close();
+        metaService.close();
     }
 
     public MetaService getSubMetaService(String schemaName) {
diff --git a/dingo-client/src/main/java/io/dingodb/client/vector/IndexService.java b/dingo-client/src/main/java/io/dingodb/client/vector/IndexService.java
@@ -86,10 +86,10 @@ public <R> R exec(
 
         Operation.Fork fork = null;
         try {
-            fork = operation.fork(Any.wrap(parameters), indexInfo);
+            fork = operation.fork(Any.wrap(parameters), indexInfo, context);
         } catch (Exception ignore) {
             indexInfo = Parameters.nonNull(cache.getIndexNewly(schemaName, indexName), "Index not found.");
-            fork = operation.fork(Any.wrap(parameters), indexInfo);
+            fork = operation.fork(Any.wrap(parameters), indexInfo, context);
         }
 
         if (operation.stateful()) {
@@ -111,7 +111,7 @@ public <R> R exec(
                 } catch (Exception e) {
                     if (OperationUtils.getCause(e) instanceof DingoClientException.InvalidRouteTableException) {
                         indexInfo = Parameters.nonNull(cache.getIndexNewly(schemaName, indexName), "Index not found.");
-                        fork = operation.fork(Any.wrap(parameters), indexInfo);
+                        fork = operation.fork(Any.wrap(parameters), indexInfo, context);
                         continue;
                     }
                     throw new RuntimeException(e);
diff --git a/dingo-client/src/main/java/io/dingodb/client/vector/Operation.java b/dingo-client/src/main/java/io/dingodb/client/vector/Operation.java
@@ -16,6 +16,7 @@
 
 package io.dingodb.client.vector;
 
+import io.dingodb.client.VectorContext;
 import io.dingodb.client.common.IndexInfo;
 import io.dingodb.sdk.common.utils.Any;
 import io.dingodb.sdk.service.entity.meta.DingoCommonId;
@@ -60,7 +61,7 @@ public <R> R result() {
         }
     }
 
-    default Fork fork(Any parameters, Index index) {
+    default Fork fork(Any parameters, Index index, VectorContext context) {
         throw new UnsupportedOperationException();
     }
 
diff --git a/dingo-client/src/main/java/io/dingodb/client/vector/VectorAddOperation.java b/dingo-client/src/main/java/io/dingodb/client/vector/VectorAddOperation.java
@@ -16,6 +16,7 @@
 
 package io.dingodb.client.vector;
 
+import io.dingodb.client.VectorContext;
 import io.dingodb.client.common.ArrayWrapperList;
 import io.dingodb.sdk.common.DingoClientException;
 import io.dingodb.sdk.common.utils.Any;
@@ -41,7 +42,8 @@ public static VectorAddOperation getInstance() {
     }
 
     @Override
-    public Fork fork(Any parameters, Index indexInfo) {
+    public Fork fork(Any parameters, Index indexInfo, VectorContext context) {
+        boolean update = context.isUpdate();
         List<VectorWithId> vectors = parameters.getValue();
         NavigableSet<Task> subTasks = new TreeSet<>(Comparator.comparing(t -> t.getRegionId().getEntityId()));
         Map<DingoCommonId, Any> subTaskMap = new HashMap<>();
@@ -62,7 +64,7 @@ public Fork fork(Any parameters, Index indexInfo) {
 
         for (int i = 0; i < vectors.size(); i++) {
             VectorWithId vector = vectors.get(i);
-            if (index.isWithAutoIncrment()) {
+            if (index.isWithAutoIncrment() && !update) {
                 long id = indexInfo.autoIncrementService.next(indexInfo.id);
                 vector.setId(id);
             }
diff --git a/dingo-client/src/main/java/io/dingodb/client/vector/VectorBatchQueryOperation.java b/dingo-client/src/main/java/io/dingodb/client/vector/VectorBatchQueryOperation.java
@@ -16,6 +16,7 @@
 
 package io.dingodb.client.vector;
 
+import io.dingodb.client.VectorContext;
 import io.dingodb.sdk.common.utils.Any;
 import io.dingodb.sdk.service.entity.common.VectorWithId;
 import io.dingodb.sdk.service.entity.index.VectorBatchQueryRequest;
@@ -43,7 +44,7 @@ public static VectorBatchQueryOperation getInstance() {
     }
 
     @Override
-    public Fork fork(Any parameters, Index indexInfo) {
+    public Fork fork(Any parameters, Index indexInfo, VectorContext context) {
         Set<Long> ids = parameters.getValue();
         NavigableSet<Task> subTasks = new TreeSet<>(Comparator.comparing(t -> t.getRegionId().getEntityId()));
         Map<DingoCommonId, Any> subTaskMap = new HashMap<>();
diff --git a/dingo-client/src/main/java/io/dingodb/client/vector/VectorCountOperation.java b/dingo-client/src/main/java/io/dingodb/client/vector/VectorCountOperation.java
@@ -16,6 +16,7 @@
 
 package io.dingodb.client.vector;
 
+import io.dingodb.client.VectorContext;
 import io.dingodb.sdk.common.utils.Any;
 import io.dingodb.sdk.service.entity.index.VectorCountRequest;
 import io.dingodb.sdk.service.entity.meta.DingoCommonId;
@@ -44,7 +45,7 @@ public boolean stateful() {
     }
 
     @Override
-    public Fork fork(Any parameters, Index indexInfo) {
+    public Fork fork(Any parameters, Index indexInfo, VectorContext context) {
         NavigableSet<Task> subTasks = new TreeSet<>(Comparator.comparing(t -> t.getRegionId().getEntityId()));
         Map<DingoCommonId, Any> subTaskMap = new HashMap<>();
 
diff --git a/dingo-client/src/main/java/io/dingodb/client/vector/VectorDeleteOperation.java b/dingo-client/src/main/java/io/dingodb/client/vector/VectorDeleteOperation.java
@@ -16,6 +16,7 @@
 
 package io.dingodb.client.vector;
 
+import io.dingodb.client.VectorContext;
 import io.dingodb.client.common.ArrayWrapperList;
 import io.dingodb.sdk.common.utils.Any;
 import io.dingodb.sdk.service.entity.index.VectorDeleteRequest;
@@ -38,7 +39,7 @@ public static VectorDeleteOperation getInstance() {
     }
 
     @Override
-    public Fork fork(Any parameters, Index indexInfo) {
+    public Fork fork(Any parameters, Index indexInfo, VectorContext context) {
         List<Long> ids = parameters.getValue();
         NavigableSet<Task> subTasks = new TreeSet<>(Comparator.comparing(t -> t.getRegionId().getEntityId()));
         Map<DingoCommonId, Any> subTaskMap = new HashMap<>();
diff --git a/dingo-client/src/main/java/io/dingodb/client/vector/VectorGetIdOperation.java b/dingo-client/src/main/java/io/dingodb/client/vector/VectorGetIdOperation.java
@@ -16,6 +16,7 @@
 
 package io.dingodb.client.vector;
 
+import io.dingodb.client.VectorContext;
 import io.dingodb.client.common.IndexInfo;
 import io.dingodb.sdk.common.utils.Any;
 import io.dingodb.sdk.service.entity.index.VectorGetBorderIdRequest;
@@ -44,7 +45,7 @@ public boolean stateful() {
     }
 
     @Override
-    public Fork fork(Any parameters, Index indexInfo) {
+    public Fork fork(Any parameters, Index indexInfo, VectorContext context) {
         Boolean isGetMin = parameters.getValue();
         NavigableSet<Task> subTasks = new TreeSet<>(Comparator.comparing(t -> t.getRegionId().getEntityId()));
         Map<DingoCommonId, Any> subTaskMap = new HashMap<>();
diff --git a/dingo-client/src/main/java/io/dingodb/client/vector/VectorGetRegionMetricsOperation.java b/dingo-client/src/main/java/io/dingodb/client/vector/VectorGetRegionMetricsOperation.java
@@ -16,6 +16,7 @@
 
 package io.dingodb.client.vector;
 
+import io.dingodb.client.VectorContext;
 import io.dingodb.sdk.common.utils.Any;
 import io.dingodb.sdk.service.entity.common.VectorIndexMetrics;
 import io.dingodb.sdk.service.entity.index.VectorGetRegionMetricsRequest;
@@ -43,7 +44,7 @@ public boolean stateful() {
     }
 
     @Override
-    public Fork fork(Any parameters, Index indexInfo) {
+    public Fork fork(Any parameters, Index indexInfo, VectorContext context) {
         NavigableSet<Task> subTasks = new TreeSet<>(Comparator.comparing(t -> t.getRegionId().getEntityId()));
         Map<DingoCommonId, Any> subTaskMap = new HashMap<>();
 
diff --git a/dingo-client/src/main/java/io/dingodb/client/vector/VectorScanQueryOperation.java b/dingo-client/src/main/java/io/dingodb/client/vector/VectorScanQueryOperation.java
@@ -16,6 +16,7 @@
 
 package io.dingodb.client.vector;
 
+import io.dingodb.client.VectorContext;
 import io.dingodb.client.common.VectorScanQuery;
 import io.dingodb.sdk.common.utils.Any;
 import io.dingodb.sdk.common.utils.Optional;
@@ -51,7 +52,7 @@ public boolean stateful() {
     }
 
     @Override
-    public Fork fork(Any parameters, Index indexInfo) {
+    public Fork fork(Any parameters, Index indexInfo, VectorContext context) {
         VectorScanQuery query = parameters.getValue();
         NavigableSet<Task> subTasks = new TreeSet<>(Comparator.comparing(t -> t.getRegionId().getEntityId()));
         Map<DingoCommonId, Any> subTaskMap = new HashMap<>();
diff --git a/dingo-client/src/main/java/io/dingodb/client/vector/VectorSearchOperation.java b/dingo-client/src/main/java/io/dingodb/client/vector/VectorSearchOperation.java
@@ -16,6 +16,7 @@
 
 package io.dingodb.client.vector;
 
+import io.dingodb.client.VectorContext;
 import io.dingodb.client.common.VectorDistanceArray;
 import io.dingodb.client.common.VectorSearch;
 import io.dingodb.sdk.common.utils.Any;
@@ -50,7 +51,7 @@ public boolean stateful() {
     }
 
     @Override
-    public Fork fork(Any parameters, Index indexInfo) {
+    public Fork fork(Any parameters, Index indexInfo, VectorContext context) {
         VectorSearch vectorSearch = parameters.getValue();
         NavigableSet<Task> subTasks = new TreeSet<>(Comparator.comparing(t -> t.getRegionId().getEntityId()));
         Map<DingoCommonId, Any> subTaskMap = new HashMap<>();
diff --git a/dingo-executor/src/main/java/io/dingodb/server/executor/schedule/LicenseCheckTask.java b/dingo-executor/src/main/java/io/dingodb/server/executor/schedule/LicenseCheckTask.java
@@ -19,6 +19,9 @@
 import io.dingodb.cluster.ClusterService;
 import io.dingodb.common.log.LogUtils;
 import io.dingodb.license.LicenseService;
+import io.dingodb.sdk.service.Services;
+import io.dingodb.sdk.service.VersionService;
+import io.dingodb.store.proxy.Configuration;
 import io.dingodb.tso.TsoService;
 import lombok.extern.slf4j.Slf4j;
 
@@ -33,14 +36,15 @@ public class LicenseCheckTask implements Runnable {
 
     @Override
     public void run() {
+        String servers = Configuration.coordinators();
         LicenseService licenseService = LicenseService.getDefault();
         if (licenseService != null) {
             long timestamp = TsoService.getDefault().timestamp();
 
             int locations = clusterService.getLocations();
             int storeMap = clusterService.getStoreMap();
 
-            boolean check = licenseService.check(timestamp, storeMap, locations);
+            boolean check = licenseService.check(timestamp, storeMap, locations, servers);
             LogUtils.info(log, "license check: {}", check);
             if (!check) {
                 clusterService.configCoordinator(true, "License check failed");
diff --git a/dingo-license-api/src/main/java/io/dingodb/license/LicenseService.java b/dingo-license-api/src/main/java/io/dingodb/license/LicenseService.java
@@ -26,6 +26,6 @@ static LicenseService getDefault() {
         );
     }
 
-    boolean check(long timestamp, int storeMap, int locations);
+    boolean check(long timestamp, int storeMap, int locations, String servers);
 
 }
diff --git a/dingo-meta-api/src/main/java/io/dingodb/meta/MetaService.java b/dingo-meta-api/src/main/java/io/dingodb/meta/MetaService.java
@@ -230,4 +230,6 @@ default TableDefinition getTableDefinition(@NonNull CommonId id) {
     void updateAutoIncrement(CommonId tableId, long autoIncrementId);
 
     long getLastId(CommonId tableId);
+
+    void close();
 }
diff --git a/dingo-store-proxy/src/main/java/io/dingodb/store/proxy/meta/MetaCache.java b/dingo-store-proxy/src/main/java/io/dingodb/store/proxy/meta/MetaCache.java
@@ -109,6 +109,8 @@ public class MetaCache {
 
     private final LoadingCache<CommonId, NavigableMap<ComparableByteArray, RangeDistribution>> distributionCache;
 
+    private boolean isClose = false;
+
     public MetaCache(Set<Location> coordinators) {
         this.coordinators = coordinators;
         this.metaService = Services.metaService(coordinators);
@@ -117,7 +119,7 @@ public MetaCache(Set<Location> coordinators) {
         this.distributionCache = buildDistributionCache();
         this.cache = new ConcurrentHashMap<>();
         Executors.execute("watch-meta", () -> {
-            while (true) {
+            while (!isClose) {
                 try {
                     watch();
                 } catch (Exception e) {
@@ -138,6 +140,11 @@ public synchronized void clear() {
         distributionCache.invalidateAll();
     }
 
+    public void close() {
+        clear();
+        isClose = true;
+    }
+
     private void watch() {
         WatchResponse response = metaService.watch(
             tso(),
@@ -146,7 +153,7 @@ private void watch() {
         clear();
         long watchId = response.getWatchId();
         long revision = -1;
-        while (true) {
+        while (!isClose) {
             response = metaService.watch(
                 tso(),
                 WatchRequest.builder().requestUnion(ProgressRequest.builder().watchId(watchId).build()).build()
diff --git a/dingo-store-proxy/src/main/java/io/dingodb/store/proxy/meta/MetaService.java b/dingo-store-proxy/src/main/java/io/dingodb/store/proxy/meta/MetaService.java
@@ -147,6 +147,12 @@ protected MetaService(DingoCommonId id, String name, io.dingodb.sdk.service.Meta
         this.cache = cache;
     }
 
+    @Override
+    public void close() {
+        api.close();
+        cache.clear();
+    }
+
     private String cleanTableName(String name) {
         return cleanName(name, "Table");
     }
diff --git a/dingo-store-proxy/src/main/java/io/dingodb/store/proxy/meta/MetaServiceApiImpl.java b/dingo-store-proxy/src/main/java/io/dingodb/store/proxy/meta/MetaServiceApiImpl.java
@@ -110,6 +110,10 @@ private MetaServiceApiImpl() {
         Executors.execute("meta-lock", this::lock, true);
     }
 
+    public void close() {
+        lockService.close();
+    }
+
     private void retryLock() {
         needLock = true;
         lockRunner.forceFollow(this::lock);
diff --git a/dingo-test/meta-local/src/main/java/io/dingodb/meta/local/LocalMetaService.java b/dingo-test/meta-local/src/main/java/io/dingodb/meta/local/LocalMetaService.java
@@ -82,6 +82,11 @@ public static void clear() {
         tableDefinitions.clear();
     }
 
+    @Override
+    public void close() {
+
+    }
+
     public static void setLocation(Location location) {
         LocalMetaService.location = location;
     }

Original file line number	Diff line number	Diff line change
`@@ -119,6 +119,7 @@ public OperationServiceV2(String coordinatorSvr) {`
`119`	`119`
`120`	`120`	`public void close() {`
`121`	`121`	`jobManager.close();`
	`122`	`+ metaService.close();`
`122`	`123`	`}`
`123`	`124`
`124`	`125`	`public MetaService getSubMetaService(String schemaName) {`
Original file line number	Diff line number	Diff line change
`@@ -16,6 +16,7 @@`
`16`	`16`
`17`	`17`	`package io.dingodb.client.vector;`
`18`	`18`
	`19`	`+import io.dingodb.client.VectorContext;`
`19`	`20`	`import io.dingodb.client.common.IndexInfo;`
`20`	`21`	`import io.dingodb.sdk.common.utils.Any;`
`21`	`22`	`import io.dingodb.sdk.service.entity.meta.DingoCommonId;`
`@@ -60,7 +61,7 @@ public <R> R result() {`
`60`	`61`	`}`
`61`	`62`	`}`
`62`	`63`
`63`		`- default Fork fork(Any parameters, Index index) {`
	`64`	`+ default Fork fork(Any parameters, Index index, VectorContext context) {`
`64`	`65`	`throw new UnsupportedOperationException();`
`65`	`66`	`}`
`66`	`67`