1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
|
// Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
package com.yahoo.vespastat;
import com.yahoo.document.BucketId;
import com.yahoo.document.BucketIdFactory;
import com.yahoo.document.DocumentId;
import com.yahoo.document.GlobalId;
import com.yahoo.document.select.BucketSelector;
import com.yahoo.document.select.BucketSet;
import com.yahoo.document.select.parser.ParseException;
import com.yahoo.documentapi.SyncParameters;
import com.yahoo.documentapi.messagebus.MessageBusDocumentAccess;
import com.yahoo.documentapi.messagebus.MessageBusSyncSession;
import com.yahoo.documentapi.messagebus.protocol.DocumentMessage;
import com.yahoo.documentapi.messagebus.protocol.GetBucketListMessage;
import com.yahoo.documentapi.messagebus.protocol.GetBucketListReply;
import com.yahoo.documentapi.messagebus.protocol.StatBucketMessage;
import com.yahoo.documentapi.messagebus.protocol.StatBucketReply;
import com.yahoo.messagebus.Reply;
import java.util.List;
/**
* This class fetches bucket information from Vespa
*
* @author bjorncs
*/
public class BucketStatsRetriever {
private final BucketIdFactory bucketIdFactory = new BucketIdFactory();
private final BucketSelector selector = new BucketSelector(bucketIdFactory);
private final MessageBusSyncSession session;
private final MessageBusDocumentAccess documentAccess;
public BucketStatsRetriever(
DocumentAccessFactory documentAccessFactory,
String route,
ShutdownHookRegistrar registrar) {
registerShutdownHook(registrar);
this.documentAccess = documentAccessFactory.createDocumentAccess();
this.session = documentAccess.createSyncSession(new SyncParameters.Builder().build());
this.session.setRoute(route);
}
private void registerShutdownHook(ShutdownHookRegistrar registrar) {
registrar.registerShutdownHook(() -> {
try {
session.destroy();
} catch (Exception e) {
// Ignore exception on shutdown
}
try {
documentAccess.shutdown();
} catch (Exception e) {
// Ignore exception on shutdown
}
});
}
public BucketId getBucketIdForType(ClientParameters.SelectionType type, String id) throws BucketStatsException {
switch (type) {
case DOCUMENT:
return bucketIdFactory.getBucketId(new DocumentId(id));
case BUCKET:
// The internal parser of BucketID is used since the Java Long.decode cannot handle unsigned longs.
return new BucketId(String.format("BucketId(%s)", id));
case GID:
return convertGidToBucketId(id);
case USER:
case GROUP:
try {
BucketSet bucketList = selector.getBucketList(createDocumentSelection(type, id));
if (bucketList.size() != 1) {
String message = String.format("Document selection must map to only one location. " +
"Specified selection matches %d locations.", bucketList.size());
throw new BucketStatsException(message);
}
return bucketList.iterator().next();
} catch (ParseException e) {
throw new BucketStatsException(String.format("Invalid id: %s (%s).", id, e.getMessage()), e);
}
default:
throw new RuntimeException("Unreachable code");
}
}
public String retrieveBucketStats(ClientParameters.SelectionType type, String id, BucketId bucketId, String bucketSpace) throws BucketStatsException {
String documentSelection = createDocumentSelection(type, id);
StatBucketMessage msg = new StatBucketMessage(bucketId, bucketSpace, documentSelection);
StatBucketReply statBucketReply = sendMessage(msg, StatBucketReply.class);
return statBucketReply.getResults();
}
public List<GetBucketListReply.BucketInfo> retrieveBucketList(BucketId bucketId, String bucketSpace) throws BucketStatsException {
GetBucketListMessage msg = new GetBucketListMessage(bucketId, bucketSpace);
GetBucketListReply bucketListReply = sendMessage(msg, GetBucketListReply.class);
return bucketListReply.getBuckets();
}
private <T extends Reply> T sendMessage(DocumentMessage msg, Class<T> expectedReply) throws BucketStatsException {
Reply reply = session.syncSend(msg);
return validateReply(reply, expectedReply);
}
private static <T extends Reply> T validateReply(Reply reply, Class<T> type) throws BucketStatsException {
if (reply.hasErrors()) {
throw new BucketStatsException(makeErrorMessage(reply));
}
if (!type.isInstance(reply)) {
throw new BucketStatsException(String.format("Unexpected reply %s: '%s'", reply.getType(), reply.toString()));
}
return type.cast(reply);
}
private static String makeErrorMessage(Reply reply) {
StringBuilder b = new StringBuilder();
b.append("Request failed: \n");
for (int i = 0; i < reply.getNumErrors(); i++) {
b.append(String.format("\t %s\n", reply.getError(i)));
}
return b.toString();
}
private static String createDocumentSelection(ClientParameters.SelectionType type, String id) {
switch (type) {
case BUCKET:
return "true";
case DOCUMENT:
return String.format("id=\"%s\"", id);
case GID:
return String.format("id.gid=\"gid(%s)\"", id);
case USER:
return String.format("id.user=%s", id);
case GROUP:
return String.format("id.group=\"%s\"", id);
default:
throw new RuntimeException("Unreachable code");
}
}
private static BucketId convertGidToBucketId(String id) throws BucketStatsException {
if (!id.matches("0x\\p{XDigit}{24}")) {
throw new BucketStatsException("Invalid gid: " + id);
}
String hexWithoutPrefix = id.substring(2);
return new GlobalId(convertHexStringToByteArray(hexWithoutPrefix)).toBucketId();
}
private static byte[] convertHexStringToByteArray(String s) {
int len = s.length();
byte[] data = new byte[len / 2];
for (int i = 0; i < len; i += 2) {
int digit1 = Character.digit(s.charAt(i), 16);
int digit2 = Character.digit(s.charAt(i + 1), 16);
data[i / 2] = (byte) ((digit1 << 4) + digit2);
}
return data;
}
public interface ShutdownHookRegistrar {
void registerShutdownHook(Runnable runnable);
}
}
|
