1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20 package org.apache.hadoop.hbase.zookeeper;
21
22 import java.io.BufferedReader;
23 import java.io.File;
24 import java.io.IOException;
25 import java.io.InputStreamReader;
26 import java.io.OutputStream;
27 import java.io.Reader;
28 import java.net.BindException;
29 import java.net.InetSocketAddress;
30 import java.net.Socket;
31 import java.util.ArrayList;
32 import java.util.List;
33 import java.util.Random;
34
35 import org.apache.commons.logging.Log;
36 import org.apache.commons.logging.LogFactory;
37 import org.apache.hadoop.conf.Configuration;
38 import org.apache.hadoop.fs.FileUtil;
39 import org.apache.hadoop.hbase.HConstants;
40 import org.apache.zookeeper.server.NIOServerCnxnFactory;
41 import org.apache.zookeeper.server.ZooKeeperServer;
42 import org.apache.zookeeper.server.persistence.FileTxnLog;
43
44
45
46
47
48
49 public class MiniZooKeeperCluster {
50 private static final Log LOG = LogFactory.getLog(MiniZooKeeperCluster.class);
51
52 private static final int TICK_TIME = 2000;
53 private static final int CONNECTION_TIMEOUT = 30000;
54
55 private boolean started;
56
57
58 private int defaultClientPort = 0;
59
60 private int clientPort;
61
62 private List<NIOServerCnxnFactory> standaloneServerFactoryList;
63 private List<ZooKeeperServer> zooKeeperServers;
64 private List<Integer> clientPortList;
65
66 private int activeZKServerIndex;
67 private int tickTime = 0;
68
69 private Configuration configuration;
70
71 public MiniZooKeeperCluster() {
72 this(new Configuration());
73 }
74
75 public MiniZooKeeperCluster(Configuration configuration) {
76 this.started = false;
77 this.configuration = configuration;
78 activeZKServerIndex = -1;
79 zooKeeperServers = new ArrayList<ZooKeeperServer>();
80 clientPortList = new ArrayList<Integer>();
81 standaloneServerFactoryList = new ArrayList<NIOServerCnxnFactory>();
82 }
83
84 public void setDefaultClientPort(int clientPort) {
85 if (clientPort <= 0) {
86 throw new IllegalArgumentException("Invalid default ZK client port: "
87 + clientPort);
88 }
89 this.defaultClientPort = clientPort;
90 }
91
92
93
94
95
96
97
98 private int selectClientPort() {
99 if (defaultClientPort > 0) {
100 return defaultClientPort;
101 }
102 return 0xc000 + new Random().nextInt(0x3f00);
103 }
104
105 public void setTickTime(int tickTime) {
106 this.tickTime = tickTime;
107 }
108
109 public int getBackupZooKeeperServerNum() {
110 return zooKeeperServers.size()-1;
111 }
112
113 public int getZooKeeperServerNum() {
114 return zooKeeperServers.size();
115 }
116
117
118 private static void setupTestEnv() {
119
120
121
122
123 System.setProperty("zookeeper.preAllocSize", "100");
124 FileTxnLog.setPreallocSize(100 * 1024);
125 }
126
127 public int startup(File baseDir) throws IOException, InterruptedException {
128 return startup(baseDir,1);
129 }
130
131
132
133
134
135
136
137
138 public int startup(File baseDir, int numZooKeeperServers) throws IOException,
139 InterruptedException {
140 if (numZooKeeperServers <= 0)
141 return -1;
142
143 setupTestEnv();
144 shutdown();
145
146 int tentativePort = selectClientPort();
147
148
149 for (int i = 0; i < numZooKeeperServers; i++) {
150 File dir = new File(baseDir, "zookeeper_"+i).getAbsoluteFile();
151 recreateDir(dir);
152 int tickTimeToUse;
153 if (this.tickTime > 0) {
154 tickTimeToUse = this.tickTime;
155 } else {
156 tickTimeToUse = TICK_TIME;
157 }
158 ZooKeeperServer server = new ZooKeeperServer(dir, dir, tickTimeToUse);
159 NIOServerCnxnFactory standaloneServerFactory;
160 while (true) {
161 try {
162 standaloneServerFactory = new NIOServerCnxnFactory();
163 standaloneServerFactory.configure(
164 new InetSocketAddress(tentativePort),
165 configuration.getInt(HConstants.ZOOKEEPER_MAX_CLIENT_CNXNS,
166 1000));
167 } catch (BindException e) {
168 LOG.debug("Failed binding ZK Server to client port: " +
169 tentativePort);
170
171 tentativePort++;
172 continue;
173 }
174 break;
175 }
176
177
178 standaloneServerFactory.startup(server);
179 if (!waitForServerUp(tentativePort, CONNECTION_TIMEOUT)) {
180 throw new IOException("Waiting for startup of standalone server");
181 }
182
183
184 clientPortList.add(tentativePort);
185 standaloneServerFactoryList.add(standaloneServerFactory);
186 zooKeeperServers.add(server);
187 }
188
189
190 activeZKServerIndex = 0;
191 started = true;
192 clientPort = clientPortList.get(activeZKServerIndex);
193 LOG.info("Started MiniZK Cluster and connect 1 ZK server " +
194 "on client port: " + clientPort);
195 return clientPort;
196 }
197
198 private void recreateDir(File dir) throws IOException {
199 if (dir.exists()) {
200 if(!FileUtil.fullyDelete(dir)) {
201 throw new IOException("Could not delete zk base directory: " + dir);
202 }
203 }
204 try {
205 dir.mkdirs();
206 } catch (SecurityException e) {
207 throw new IOException("creating dir: " + dir, e);
208 }
209 }
210
211
212
213
214 public void shutdown() throws IOException {
215 if (!started) {
216 return;
217 }
218
219
220 for (int i = 0; i < standaloneServerFactoryList.size(); i++) {
221 NIOServerCnxnFactory standaloneServerFactory =
222 standaloneServerFactoryList.get(i);
223 int clientPort = clientPortList.get(i);
224
225 standaloneServerFactory.shutdown();
226 if (!waitForServerDown(clientPort, CONNECTION_TIMEOUT)) {
227 throw new IOException("Waiting for shutdown of standalone server");
228 }
229 }
230 for (ZooKeeperServer zkServer: zooKeeperServers) {
231
232 zkServer.getZKDatabase().close();
233 }
234
235
236 started = false;
237 activeZKServerIndex = 0;
238 standaloneServerFactoryList.clear();
239 clientPortList.clear();
240 zooKeeperServers.clear();
241
242 LOG.info("Shutdown MiniZK cluster with all ZK servers");
243 }
244
245
246
247
248
249
250 public int killCurrentActiveZooKeeperServer() throws IOException,
251 InterruptedException {
252 if (!started || activeZKServerIndex < 0 ) {
253 return -1;
254 }
255
256
257 NIOServerCnxnFactory standaloneServerFactory =
258 standaloneServerFactoryList.get(activeZKServerIndex);
259 int clientPort = clientPortList.get(activeZKServerIndex);
260
261 standaloneServerFactory.shutdown();
262 if (!waitForServerDown(clientPort, CONNECTION_TIMEOUT)) {
263 throw new IOException("Waiting for shutdown of standalone server");
264 }
265
266
267 standaloneServerFactoryList.remove(activeZKServerIndex);
268 clientPortList.remove(activeZKServerIndex);
269 zooKeeperServers.remove(activeZKServerIndex);
270 LOG.info("Kill the current active ZK servers in the cluster " +
271 "on client port: " + clientPort);
272
273 if (standaloneServerFactoryList.size() == 0) {
274
275 return -1;
276 }
277 clientPort = clientPortList.get(activeZKServerIndex);
278 LOG.info("Activate a backup zk server in the cluster " +
279 "on client port: " + clientPort);
280
281 return clientPort;
282 }
283
284
285
286
287
288
289 public void killOneBackupZooKeeperServer() throws IOException,
290 InterruptedException {
291 if (!started || activeZKServerIndex < 0 ||
292 standaloneServerFactoryList.size() <= 1) {
293 return ;
294 }
295
296 int backupZKServerIndex = activeZKServerIndex+1;
297
298 NIOServerCnxnFactory standaloneServerFactory =
299 standaloneServerFactoryList.get(backupZKServerIndex);
300 int clientPort = clientPortList.get(backupZKServerIndex);
301
302 standaloneServerFactory.shutdown();
303 if (!waitForServerDown(clientPort, CONNECTION_TIMEOUT)) {
304 throw new IOException("Waiting for shutdown of standalone server");
305 }
306
307
308 standaloneServerFactoryList.remove(backupZKServerIndex);
309 clientPortList.remove(backupZKServerIndex);
310 zooKeeperServers.remove(backupZKServerIndex);
311 LOG.info("Kill one backup ZK servers in the cluster " +
312 "on client port: " + clientPort);
313 }
314
315
316 private static boolean waitForServerDown(int port, long timeout) {
317 long start = System.currentTimeMillis();
318 while (true) {
319 try {
320 Socket sock = new Socket("localhost", port);
321 try {
322 OutputStream outstream = sock.getOutputStream();
323 outstream.write("stat".getBytes());
324 outstream.flush();
325 } finally {
326 sock.close();
327 }
328 } catch (IOException e) {
329 return true;
330 }
331
332 if (System.currentTimeMillis() > start + timeout) {
333 break;
334 }
335 try {
336 Thread.sleep(250);
337 } catch (InterruptedException e) {
338
339 }
340 }
341 return false;
342 }
343
344
345 private static boolean waitForServerUp(int port, long timeout) {
346 long start = System.currentTimeMillis();
347 while (true) {
348 try {
349 Socket sock = new Socket("localhost", port);
350 BufferedReader reader = null;
351 try {
352 OutputStream outstream = sock.getOutputStream();
353 outstream.write("stat".getBytes());
354 outstream.flush();
355
356 Reader isr = new InputStreamReader(sock.getInputStream());
357 reader = new BufferedReader(isr);
358 String line = reader.readLine();
359 if (line != null && line.startsWith("Zookeeper version:")) {
360 return true;
361 }
362 } finally {
363 sock.close();
364 if (reader != null) {
365 reader.close();
366 }
367 }
368 } catch (IOException e) {
369
370 LOG.info("server localhost:" + port + " not up " + e);
371 }
372
373 if (System.currentTimeMillis() > start + timeout) {
374 break;
375 }
376 try {
377 Thread.sleep(250);
378 } catch (InterruptedException e) {
379
380 }
381 }
382 return false;
383 }
384
385 public int getClientPort() {
386 return clientPort;
387 }
388 }