Skip to content

Commit 425be21

Browse files
yhmoyhmo
authored
Add flushAll/getFlushAllState for MilvusClientV2 (#1862)
Signed-off-by: yhmo <yihua.mo@zilliz.com>
1 parent 0c9c715 commit 425be21

7 files changed

Lines changed: 372 additions & 1 deletion

File tree

‎sdk-core/src/main/java/io/milvus/v2/client/MilvusClientV2.java‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1267,6 +1267,34 @@ public void flush(FlushReq request) {
12671267
utilityService.waitFlush(tempBlockingStub, response);
12681268
}
12691269

1270+
/**
1271+
* trigger a flush action for all collections in server side
1272+
*
1273+
* @param request flush all request
1274+
* @return FlushAllResp
1275+
*/
1276+
public FlushAllResp flushAll(FlushAllReq request) {
1277+
FlushAllResp response = rpcUtils.retry(() -> utilityService.flushAll(this.getRpcStub(), request));
1278+
1279+
MilvusServiceGrpc.MilvusServiceBlockingStub tempBlockingStub =
1280+
MilvusServiceGrpc.newBlockingStub(channel).withWaitForReady();
1281+
if (request.getWaitFlushedTimeoutMs() > 0L) {
1282+
tempBlockingStub = tempBlockingStub.withDeadlineAfter(request.getWaitFlushedTimeoutMs(), TimeUnit.MILLISECONDS);
1283+
}
1284+
utilityService.waitFlushAll(tempBlockingStub, response, request);
1285+
return response;
1286+
}
1287+
1288+
/**
1289+
* Gets the flush all state.
1290+
*
1291+
* @param request get flush all state request
1292+
* @return GetFlushAllStateResp
1293+
*/
1294+
public GetFlushAllStateResp getFlushAllState(GetFlushAllStateReq request) {
1295+
return rpcUtils.retry(() -> utilityService.getFlushAllState(this.getRpcStub(), request));
1296+
}
1297+
12701298
/**
12711299
* Gets the information of persistent segments from data node, including row count,
12721300
* persistence state(growing or flushed), etc.

‎sdk-core/src/main/java/io/milvus/v2/service/utility/UtilityService.java‎

Lines changed: 66 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,40 @@ public FlushResp flush(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub,
9090
.build();
9191
}
9292

93+
public FlushAllResp flushAll(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub, FlushAllReq request) {
94+
String dbName = request.getDatabaseName();
95+
String title = String.format("Flush all in database: '%s'", dbName);
96+
97+
FlushAllRequest.Builder builder = FlushAllRequest.newBuilder();
98+
if (StringUtils.isNotEmpty(dbName)) {
99+
builder.setDbName(dbName);
100+
}
101+
FlushAllResponse response = blockingStub.flushAll(builder.build());
102+
rpcUtils.handleResponse(title, response.getStatus());
103+
104+
return FlushAllResp.builder()
105+
.flushAllTs(response.getFlushAllTs())
106+
.build();
107+
}
108+
109+
public GetFlushAllStateResp getFlushAllState(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub,
110+
GetFlushAllStateReq request) {
111+
String dbName = request.getDatabaseName();
112+
String title = String.format("Get flush all state in database: '%s'", dbName);
113+
114+
GetFlushAllStateRequest.Builder builder = GetFlushAllStateRequest.newBuilder()
115+
.setFlushAllTs(request.getFlushAllTs());
116+
if (StringUtils.isNotEmpty(dbName)) {
117+
builder.setDbName(dbName);
118+
}
119+
GetFlushAllStateResponse response = blockingStub.getFlushAllState(builder.build());
120+
rpcUtils.handleResponse(title, response.getStatus());
121+
122+
return GetFlushAllStateResp.builder()
123+
.flushed(response.getFlushed())
124+
.build();
125+
}
126+
93127
// this method is internal use, not expose to user
94128
public Void waitFlush(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub, FlushResp flushResp) {
95129
Map<String, List<Long>> collectionSegmentIDs = flushResp.getCollectionSegmentIDs();
@@ -109,7 +143,8 @@ public Void waitFlush(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub,
109143
try {
110144
TimeUnit.SECONDS.sleep(1);
111145
} catch (InterruptedException t) {
112-
System.out.println("Interrupted: " + t.getMessage());
146+
Thread.currentThread().interrupt();
147+
logger.warn("Interrupted while waiting for flush", t);
113148
break;
114149
}
115150
}
@@ -119,6 +154,36 @@ public Void waitFlush(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub,
119154
return null;
120155
}
121156

157+
// this method is internal use, not expose to user
158+
public Void waitFlushAll(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub, FlushAllResp flushAllResp,
159+
FlushAllReq request) {
160+
boolean flushed = false;
161+
long start = System.currentTimeMillis();
162+
while (!flushed) {
163+
GetFlushAllStateResp flushAllStateResp = getFlushAllState(blockingStub, GetFlushAllStateReq.builder()
164+
.databaseName(request.getDatabaseName())
165+
.flushAllTs(flushAllResp.getFlushAllTs())
166+
.build());
167+
168+
flushed = Boolean.TRUE.equals(flushAllStateResp.getFlushed());
169+
if (!flushed) {
170+
Long timeout = request.getWaitFlushedTimeoutMs();
171+
if (timeout != null && timeout > 0L && System.currentTimeMillis() - start > timeout) {
172+
throw new MilvusClientException(ErrorCode.CLIENT_ERROR,
173+
"wait for flush all timeout, flush_all_ts: " + flushAllResp.getFlushAllTs());
174+
}
175+
try {
176+
TimeUnit.SECONDS.sleep(5);
177+
} catch (InterruptedException t) {
178+
Thread.currentThread().interrupt();
179+
throw new MilvusClientException(ErrorCode.CLIENT_ERROR, "Interrupted while waiting for flush all");
180+
}
181+
}
182+
}
183+
184+
return null;
185+
}
186+
122187
public CompactResp compact(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub, CompactReq request) {
123188
String dbName = request.getDatabaseName();
124189
String collectionName = request.getCollectionName();

‎sdk-core/src/main/java/io/milvus/v2/service/utility/request/FlushAllReq.java‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package io.milvus.v2.service.utility.request;
21+
22+
public class FlushAllReq {
23+
private String databaseName;
24+
private Long waitFlushedTimeoutMs;
25+
26+
private FlushAllReq(FlushAllReqBuilder builder) {
27+
this.databaseName = builder.databaseName;
28+
this.waitFlushedTimeoutMs = builder.waitFlushedTimeoutMs;
29+
}
30+
31+
public static FlushAllReqBuilder builder() {
32+
return new FlushAllReqBuilder();
33+
}
34+
35+
public String getDatabaseName() {
36+
return databaseName;
37+
}
38+
39+
public void setDatabaseName(String databaseName) {
40+
this.databaseName = databaseName;
41+
}
42+
43+
public Long getWaitFlushedTimeoutMs() {
44+
return waitFlushedTimeoutMs;
45+
}
46+
47+
public void setWaitFlushedTimeoutMs(Long waitFlushedTimeoutMs) {
48+
this.waitFlushedTimeoutMs = waitFlushedTimeoutMs;
49+
}
50+
51+
@Override
52+
public String toString() {
53+
return "FlushAllReq{" +
54+
"databaseName='" + databaseName + '\'' +
55+
", waitFlushedTimeoutMs=" + waitFlushedTimeoutMs +
56+
'}';
57+
}
58+
59+
public static class FlushAllReqBuilder {
60+
private String databaseName;
61+
private Long waitFlushedTimeoutMs = 0L;
62+
63+
public FlushAllReqBuilder databaseName(String databaseName) {
64+
this.databaseName = databaseName;
65+
return this;
66+
}
67+
68+
public FlushAllReqBuilder waitFlushedTimeoutMs(Long waitFlushedTimeoutMs) {
69+
this.waitFlushedTimeoutMs = waitFlushedTimeoutMs;
70+
return this;
71+
}
72+
73+
public FlushAllReq build() {
74+
return new FlushAllReq(this);
75+
}
76+
}
77+
}

‎sdk-core/src/main/java/io/milvus/v2/service/utility/request/GetFlushAllStateReq.java‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package io.milvus.v2.service.utility.request;
21+
22+
public class GetFlushAllStateReq {
23+
private String databaseName;
24+
private Long flushAllTs;
25+
26+
private GetFlushAllStateReq(GetFlushAllStateReqBuilder builder) {
27+
this.databaseName = builder.databaseName;
28+
this.flushAllTs = builder.flushAllTs;
29+
}
30+
31+
public static GetFlushAllStateReqBuilder builder() {
32+
return new GetFlushAllStateReqBuilder();
33+
}
34+
35+
public String getDatabaseName() {
36+
return databaseName;
37+
}
38+
39+
public void setDatabaseName(String databaseName) {
40+
this.databaseName = databaseName;
41+
}
42+
43+
public Long getFlushAllTs() {
44+
return flushAllTs;
45+
}
46+
47+
public void setFlushAllTs(Long flushAllTs) {
48+
this.flushAllTs = flushAllTs;
49+
}
50+
51+
@Override
52+
public String toString() {
53+
return "GetFlushAllStateReq{" +
54+
"databaseName='" + databaseName + '\'' +
55+
", flushAllTs=" + flushAllTs +
56+
'}';
57+
}
58+
59+
public static class GetFlushAllStateReqBuilder {
60+
private String databaseName;
61+
private Long flushAllTs = 0L;
62+
63+
public GetFlushAllStateReqBuilder databaseName(String databaseName) {
64+
this.databaseName = databaseName;
65+
return this;
66+
}
67+
68+
public GetFlushAllStateReqBuilder flushAllTs(Long flushAllTs) {
69+
this.flushAllTs = flushAllTs;
70+
return this;
71+
}
72+
73+
public GetFlushAllStateReq build() {
74+
return new GetFlushAllStateReq(this);
75+
}
76+
}
77+
}

‎sdk-core/src/main/java/io/milvus/v2/service/utility/response/FlushAllResp.java‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package io.milvus.v2.service.utility.response;
21+
22+
public class FlushAllResp {
23+
private Long flushAllTs;
24+
25+
private FlushAllResp(FlushAllRespBuilder builder) {
26+
this.flushAllTs = builder.flushAllTs;
27+
}
28+
29+
public static FlushAllRespBuilder builder() {
30+
return new FlushAllRespBuilder();
31+
}
32+
33+
public Long getFlushAllTs() {
34+
return flushAllTs;
35+
}
36+
37+
public void setFlushAllTs(Long flushAllTs) {
38+
this.flushAllTs = flushAllTs;
39+
}
40+
41+
@Override
42+
public String toString() {
43+
return "FlushAllResp{" +
44+
"flushAllTs=" + flushAllTs +
45+
'}';
46+
}
47+
48+
public static class FlushAllRespBuilder {
49+
private Long flushAllTs = 0L;
50+
51+
public FlushAllRespBuilder flushAllTs(Long flushAllTs) {
52+
this.flushAllTs = flushAllTs;
53+
return this;
54+
}
55+
56+
public FlushAllResp build() {
57+
return new FlushAllResp(this);
58+
}
59+
}
60+
}

0 commit comments

Comments
 (0)