Adding an assert, plus a conversion of malloc to calloc
[awesomized/libmemcached] / clients / memslap.c
1 #include "libmemcached/common.h"
2 #include <stdio.h>
3 #include <stdlib.h>
4 #include <string.h>
5 #include <sys/types.h>
6 #include <sys/stat.h>
7 #include <sys/types.h>
8 #include <sys/mman.h>
9 #include <fcntl.h>
10 #include <sys/time.h>
11 #include <getopt.h>
12 #include <pthread.h>
13 #include <assert.h>
14
15 #include <libmemcached/memcached.h>
16
17 #include "client_options.h"
18 #include "utilities.h"
19 #include "generator.h"
20 #include "execute.h"
21
22 #define DEFAULT_INITIAL_LOAD 10000
23 #define DEFAULT_EXECUTE_NUMBER 10000
24 #define DEFAULT_CONCURRENCY 1
25
26 #define PROGRAM_NAME "memslap"
27 #define PROGRAM_DESCRIPTION "Generates a load against a memcached custer of servers."
28
29 /* Global Thread counter */
30 volatile unsigned int thread_counter;
31 pthread_mutex_t counter_mutex;
32 pthread_cond_t count_threshhold;
33 volatile unsigned int master_wakeup;
34 pthread_mutex_t sleeper_mutex;
35 pthread_cond_t sleep_threshhold;
36
37 void *run_task(void *p);
38
39 /* Types */
40 typedef struct conclusions_st conclusions_st;
41 typedef struct thread_context_st thread_context_st;
42 typedef enum {
43 SET_TEST,
44 GET_TEST,
45 } test_type;
46
47 struct thread_context_st {
48 unsigned int key_count;
49 pairs_st *initial_pairs;
50 unsigned int initial_number;
51 pairs_st *execute_pairs;
52 unsigned int execute_number;
53 test_type test;
54 memcached_st *memc;
55 };
56
57 struct conclusions_st {
58 long int load_time;
59 long int read_time;
60 unsigned int rows_loaded;
61 unsigned int rows_read;
62 };
63
64 /* Prototypes */
65 void options_parse(int argc, char *argv[]);
66 void conclusions_print(conclusions_st *conclusion);
67 void scheduler(memcached_server_st *servers, conclusions_st *conclusion);
68 pairs_st *load_create_data(memcached_st *memc, unsigned int number_of,
69 unsigned int *actual_loaded);
70 void flush_all(memcached_st *memc);
71
72 static int opt_binary= 0;
73 static int opt_verbose= 0;
74 static int opt_flush= 0;
75 static int opt_non_blocking_io= 0;
76 static int opt_tcp_nodelay= 0;
77 static unsigned int opt_execute_number= 0;
78 static unsigned int opt_createial_load= 0;
79 static unsigned int opt_concurrency= 0;
80 static int opt_displayflag= 0;
81 static char *opt_servers= NULL;
82 static int opt_udp_io= 0;
83 test_type opt_test= SET_TEST;
84
85 int main(int argc, char *argv[])
86 {
87 conclusions_st conclusion;
88 memcached_server_st *servers;
89
90 memset(&conclusion, 0, sizeof(conclusions_st));
91
92 srandom(time(NULL));
93 options_parse(argc, argv);
94
95 if (!opt_servers)
96 {
97 char *temp;
98
99 if ((temp= getenv("MEMCACHED_SERVERS")))
100 opt_servers= strdup(temp);
101 else
102 {
103 fprintf(stderr, "No Servers provided\n");
104 exit(1);
105 }
106 }
107
108 servers= memcached_servers_parse(opt_servers);
109
110 pthread_mutex_init(&counter_mutex, NULL);
111 pthread_cond_init(&count_threshhold, NULL);
112 pthread_mutex_init(&sleeper_mutex, NULL);
113 pthread_cond_init(&sleep_threshhold, NULL);
114
115 scheduler(servers, &conclusion);
116
117 free(opt_servers);
118
119 (void)pthread_mutex_destroy(&counter_mutex);
120 (void)pthread_cond_destroy(&count_threshhold);
121 (void)pthread_mutex_destroy(&sleeper_mutex);
122 (void)pthread_cond_destroy(&sleep_threshhold);
123 conclusions_print(&conclusion);
124 memcached_server_list_free(servers);
125
126 return 0;
127 }
128
129 void scheduler(memcached_server_st *servers, conclusions_st *conclusion)
130 {
131 unsigned int x;
132 unsigned int actual_loaded= 0; /* Fix warning */
133 memcached_st *memc;
134
135 struct timeval start_time, end_time;
136 pthread_t mainthread; /* Thread descriptor */
137 pthread_attr_t attr; /* Thread attributes */
138 pairs_st *pairs= NULL;
139
140 pthread_attr_init(&attr);
141 pthread_attr_setdetachstate(&attr,
142 PTHREAD_CREATE_DETACHED);
143
144 memc= memcached_create(NULL);
145
146 /* We need to set udp behavior before adding servers to the client */
147 if (opt_udp_io)
148 {
149 memcached_behavior_set(memc, MEMCACHED_BEHAVIOR_USE_UDP, opt_udp_io);
150 unsigned int x= 0;
151 for(x= 0; x < servers[0].count; x++ )
152 servers[x].type= MEMCACHED_CONNECTION_UDP;
153 }
154 memcached_server_push(memc, servers);
155
156 memcached_behavior_set(memc, MEMCACHED_BEHAVIOR_BINARY_PROTOCOL, opt_binary);
157
158 if (opt_flush)
159 flush_all(memc);
160 if (opt_createial_load)
161 pairs= load_create_data(memc, opt_createial_load, &actual_loaded);
162
163 /* We set this after we have loaded */
164 {
165 if (opt_non_blocking_io)
166 memcached_behavior_set(memc, MEMCACHED_BEHAVIOR_NO_BLOCK, 1);
167 if (opt_tcp_nodelay)
168 memcached_behavior_set(memc, MEMCACHED_BEHAVIOR_TCP_NODELAY, 1);
169 }
170
171
172 pthread_mutex_lock(&counter_mutex);
173 thread_counter= 0;
174
175 pthread_mutex_lock(&sleeper_mutex);
176 master_wakeup= 1;
177 pthread_mutex_unlock(&sleeper_mutex);
178
179 for (x= 0; x < opt_concurrency; x++)
180 {
181 thread_context_st *context;
182 context= (thread_context_st *)calloc(1, sizeof(thread_context_st));
183
184 context->memc= memcached_clone(NULL, memc);
185 context->test= opt_test;
186
187 context->initial_pairs= pairs;
188 context->initial_number= actual_loaded;
189
190 if (opt_test == SET_TEST)
191 {
192 context->execute_pairs= pairs_generate(opt_execute_number, 400);
193 context->execute_number= opt_execute_number;
194 }
195
196 /* now you create the thread */
197 if (pthread_create(&mainthread, &attr, run_task,
198 (void *)context) != 0)
199 {
200 fprintf(stderr,"Could not create thread\n");
201 exit(1);
202 }
203 thread_counter++;
204 }
205
206 pthread_mutex_unlock(&counter_mutex);
207 pthread_attr_destroy(&attr);
208
209 pthread_mutex_lock(&sleeper_mutex);
210 master_wakeup= 0;
211 pthread_mutex_unlock(&sleeper_mutex);
212 pthread_cond_broadcast(&sleep_threshhold);
213
214 gettimeofday(&start_time, NULL);
215 /*
216 We loop until we know that all children have cleaned up.
217 */
218 pthread_mutex_lock(&counter_mutex);
219 while (thread_counter)
220 pthread_cond_wait(&count_threshhold, &counter_mutex);
221 pthread_mutex_unlock(&counter_mutex);
222
223 gettimeofday(&end_time, NULL);
224
225 conclusion->load_time= timedif(end_time, start_time);
226 conclusion->read_time= timedif(end_time, start_time);
227 pairs_free(pairs);
228 memcached_free(memc);
229 }
230
231 void options_parse(int argc, char *argv[])
232 {
233 memcached_programs_help_st help_options[]=
234 {
235 {0},
236 };
237
238 static struct option long_options[]=
239 {
240 {"concurrency", required_argument, NULL, OPT_SLAP_CONCURRENCY},
241 {"debug", no_argument, &opt_verbose, OPT_DEBUG},
242 {"execute-number", required_argument, NULL, OPT_SLAP_EXECUTE_NUMBER},
243 {"flag", no_argument, &opt_displayflag, OPT_FLAG},
244 {"flush", no_argument, &opt_flush, OPT_FLUSH},
245 {"help", no_argument, NULL, OPT_HELP},
246 {"initial-load", required_argument, NULL, OPT_SLAP_INITIAL_LOAD}, /* Number to load initially */
247 {"non-blocking", no_argument, &opt_non_blocking_io, OPT_SLAP_NON_BLOCK},
248 {"servers", required_argument, NULL, OPT_SERVERS},
249 {"tcp-nodelay", no_argument, &opt_tcp_nodelay, OPT_SLAP_TCP_NODELAY},
250 {"test", required_argument, NULL, OPT_SLAP_TEST},
251 {"verbose", no_argument, &opt_verbose, OPT_VERBOSE},
252 {"version", no_argument, NULL, OPT_VERSION},
253 {"binary", no_argument, NULL, OPT_BINARY},
254 {"udp", no_argument, NULL, OPT_UDP},
255 {0, 0, 0, 0},
256 };
257
258 int option_index= 0;
259 int option_rv;
260
261 while (1)
262 {
263 option_rv= getopt_long(argc, argv, "Vhvds:", long_options, &option_index);
264 if (option_rv == -1) break;
265 switch (option_rv)
266 {
267 case 0:
268 break;
269 case OPT_UDP:
270 if (opt_test == GET_TEST)
271 {
272 fprintf(stderr, "You can not run a get test in UDP mode. UDP mode "
273 "does not currently support get ops.\n");
274 exit(1);
275 }
276 opt_udp_io= 1;
277 break;
278 case OPT_BINARY:
279 opt_binary = 1;
280 break;
281 case OPT_VERBOSE: /* --verbose or -v */
282 opt_verbose = OPT_VERBOSE;
283 break;
284 case OPT_DEBUG: /* --debug or -d */
285 opt_verbose = OPT_DEBUG;
286 break;
287 case OPT_VERSION: /* --version or -V */
288 version_command(PROGRAM_NAME);
289 break;
290 case OPT_HELP: /* --help or -h */
291 help_command(PROGRAM_NAME, PROGRAM_DESCRIPTION, long_options, help_options);
292 break;
293 case OPT_SERVERS: /* --servers or -s */
294 opt_servers= strdup(optarg);
295 break;
296 case OPT_SLAP_TEST:
297 if (!strcmp(optarg, "get"))
298 {
299 if (opt_udp_io == 1)
300 {
301 fprintf(stderr, "You can not run a get test in UDP mode. UDP mode "
302 "does not currently support get ops.\n");
303 exit(1);
304 }
305 opt_test= GET_TEST ;
306 }
307 else if (!strcmp(optarg, "set"))
308 opt_test= SET_TEST;
309 else
310 {
311 fprintf(stderr, "Your test, %s, is not a known test\n", optarg);
312 exit(1);
313 }
314 break;
315 case OPT_SLAP_CONCURRENCY:
316 opt_concurrency= strtol(optarg, (char **)NULL, 10);
317 break;
318 case OPT_SLAP_EXECUTE_NUMBER:
319 opt_execute_number= strtol(optarg, (char **)NULL, 10);
320 break;
321 case OPT_SLAP_INITIAL_LOAD:
322 opt_createial_load= strtol(optarg, (char **)NULL, 10);
323 break;
324 case '?':
325 /* getopt_long already printed an error message. */
326 exit(1);
327 default:
328 abort();
329 }
330 }
331
332 if (opt_test == GET_TEST && opt_createial_load == 0)
333 opt_createial_load= DEFAULT_INITIAL_LOAD;
334
335 if (opt_execute_number == 0)
336 opt_execute_number= DEFAULT_EXECUTE_NUMBER;
337
338 if (opt_concurrency == 0)
339 opt_concurrency= DEFAULT_CONCURRENCY;
340 }
341
342 void conclusions_print(conclusions_st *conclusion)
343 {
344 printf("\tThreads connecting to servers %u\n", opt_concurrency);
345 #ifdef NOT_FINISHED
346 printf("\tLoaded %u rows\n", conclusion->rows_loaded);
347 printf("\tRead %u rows\n", conclusion->rows_read);
348 #endif
349 if (opt_test == SET_TEST)
350 printf("\tTook %ld.%03ld seconds to load data\n", conclusion->load_time / 1000,
351 conclusion->load_time % 1000);
352 else
353 printf("\tTook %ld.%03ld seconds to read data\n", conclusion->read_time / 1000,
354 conclusion->read_time % 1000);
355 }
356
357 void *run_task(void *p)
358 {
359 thread_context_st *context= (thread_context_st *)p;
360 memcached_st *memc;
361
362 memc= context->memc;
363
364 pthread_mutex_lock(&sleeper_mutex);
365 while (master_wakeup)
366 {
367 pthread_cond_wait(&sleep_threshhold, &sleeper_mutex);
368 }
369 pthread_mutex_unlock(&sleeper_mutex);
370
371 /* Do Stuff */
372 switch (context->test)
373 {
374 case SET_TEST:
375 assert(context->execute_pairs);
376 execute_set(memc, context->execute_pairs, context->execute_number);
377 break;
378 case GET_TEST:
379 execute_get(memc, context->initial_pairs, context->initial_number);
380 break;
381 }
382
383 memcached_free(memc);
384
385 if (context->execute_pairs)
386 pairs_free(context->execute_pairs);
387
388 free(context);
389
390 pthread_mutex_lock(&counter_mutex);
391 thread_counter--;
392 pthread_cond_signal(&count_threshhold);
393 pthread_mutex_unlock(&counter_mutex);
394
395 return NULL;
396 }
397
398 void flush_all(memcached_st *memc)
399 {
400 memcached_flush(memc, 0);
401 }
402
403 pairs_st *load_create_data(memcached_st *memc, unsigned int number_of,
404 unsigned int *actual_loaded)
405 {
406 memcached_st *clone;
407 pairs_st *pairs;
408
409 clone= memcached_clone(NULL, memc);
410 /* We always used non-blocking IO for load since it is faster */
411 memcached_behavior_set(clone, MEMCACHED_BEHAVIOR_NO_BLOCK, 0);
412
413 pairs= pairs_generate(number_of, 400);
414 *actual_loaded= execute_set(clone, pairs, number_of);
415
416 memcached_free(clone);
417
418 return pairs;
419 }