1 #include <io.h>
2 #include <byte.h>
3 #include <str.h>
4 #include <fmt.h>
5 #include <scan.h>
6 #include <socket.h>
7 #include <errmsg.h>
8 #include <dns.h>
9 #include <ip6.h>
10 #include <textcode.h>
11 #include <signal.h>
12 #include <unistd.h>
13 #include <errno.h>
14 #include <buffer.h>
15 #include <fcntl.h>
16 #include <string.h>
17 #include "havealloca.h"
18 
usage()19 void usage() {
20   die(0,"usage: bench [-n requests] [-c concurrency] [-t timeout] [-k] [-K count]\n"
21         "       [-C cookie-file] [-p] ([http://]host[:port]/uri|@host:port)");
22 }
23 
24 unsigned long r[10];
25 unsigned long kaputt;
26 
make_connection(char * ip,uint16 port,uint32 scope_id,int s)27 static int make_connection(char* ip,uint16 port,uint32 scope_id,int s) {
28   int v6=byte_diff(ip,12,V4mappedprefix);
29   if (v6) {
30     if (s==-1) {
31       s=socket_tcp6();
32       if (s==-1) return -1;
33     }
34     if (socket_connect6(s,ip,port,scope_id)==-1) {
35       if (errno==EAGAIN || errno==EINPROGRESS || errno==EALREADY || errno==EISCONN)
36 	return s;
37       ++kaputt;
38       if (errno!=ECONNREFUSED && errno!=ECONNRESET)
39 	carpsys("socket_connect6");
40       close(s);
41       return -1;
42     }
43   } else {
44     if (s==-1) {
45       s=socket_tcp4();
46       if (s==-1) return -1;
47     }
48     if (socket_connect4(s,ip+12,port)==-1) {
49       if (errno==EAGAIN || errno==EINPROGRESS || errno==EALREADY || errno==EISCONN)
50 	return s;
51       ++kaputt;
52       if (errno!=ECONNREFUSED && errno!=ECONNRESET)
53 	carpsys("socket_connect6");
54       close(s);
55       return -1;
56     }
57   }
58   return s;
59 }
60 
61 buffer* cookies;
62 
cookiefile(const char * s)63 void cookiefile(const char* s) {
64   static buffer cookiebuffer;
65   static char cookiebuf[8192];
66   int fd;
67   if (!s) {
68     lseek(cookiebuffer.fd,0,SEEK_SET);
69     buffer_init(&cookiebuffer,read,cookiebuffer.fd,cookiebuf,sizeof cookiebuf);
70   } else {
71     fd=open(s,O_RDONLY);
72     if (fd==-1) die(1,"could not open cookie file \"",s,"\"!");
73     if (cookiebuffer.fd!=0) close(cookiebuffer.fd);
74     buffer_init(&cookiebuffer,read,fd,cookiebuf,sizeof cookiebuf);
75     cookies=&cookiebuffer;
76   }
77 }
78 
nextcookie(char * dest,unsigned long destlen)79 int nextcookie(char* dest,unsigned long destlen) {
80   int len;
81   if (!cookies) return -1;
82   if ((len=buffer_getline(cookies,dest,destlen))) {
83     if (dest[len]!='\n')
84       die(0,"line too long: ",dest);
85     dest[len]=0;
86   } else {
87     cookiefile(0);
88     if ((len=buffer_getline(cookies,dest,destlen))) {
89       if (dest[len]!='\n')
90 	die(0,"line too long: ",dest);
91       dest[len]=0;
92     } else
93       return -1;
94   }
95   return len;
96 }
97 
main(int argc,char * argv[])98 int main(int argc,char* argv[]) {
99   char server[1024];
100   int* fds;
101   int* avail;
102   int* keepleft;
103   size_t* requests;
104   long long* expected;
105   unsigned long n=10000;	/* requests */
106   unsigned long c=10;		/* concurrency */
107   unsigned long t=0;		/* time limit in seconds */
108   unsigned long k=0;		/* keep-alive */
109   unsigned long K=1;		/* keep-alive counter */
110   int report=0;
111   unsigned long long errors=0;
112   unsigned long long bytes=0;
113   int v=0;
114   int p=0;	// print profile
115   unsigned long i,done;
116   uint16 port=80;
117   uint32 scope_id=0;
118   stralloc ips={0};
119   char* request,* krequest;
120   unsigned long rlen, krlen;
121   tai6464 first,now,next,last;
122   enum { SAME, REPLAY } mode;
123   char* hostname;
124 
125   server[0]=0;
126 
127   errmsg_iam("bench");
128 #ifndef __MINGW32__
129   signal(SIGPIPE,SIG_IGN);
130 #endif
131 
132   for (;;) {
133     int i;
134     int ch=getopt(argc,argv,"n:c:t:kvK:C:rp");
135     if (ch==-1) break;
136     switch (ch) {
137     case 'r':
138       report=1;
139       break;
140     case 'n':
141       i=scan_ulong(optarg,&n);
142       if (i==0) die(1,"could not parse -n argument \"",optarg,"\".\n");
143       break;
144     case 'c':
145       i=scan_ulong(optarg,&c);
146       if (i==0) die(1,"could not parse -c argument \"",optarg,"\".\n");
147       break;
148     case 't':
149       i=scan_ulong(optarg,&t);
150       if (i==0) die(1,"could not parse -t argument \"",optarg,"\".\n");
151       break;
152     case 'k':
153       k=1;
154       break;
155     case 'K':
156       i=scan_ulong(optarg,&K);
157       break;
158     case 'v':
159       v=1;
160       break;
161     case 'C':
162       cookiefile(optarg);
163       break;
164     case 'p':
165       p=1;
166     case '?':
167       break;
168     default:
169       usage();
170     }
171   }
172   if (n<1 || c<1 || !argv[optind]) usage();
173   if (c>100000) {
174     carp("concurrency limited to 100000 in this build.");
175     c=100000;
176   }
177 
178   if (argv[optind][0]=='@') {
179     mode=REPLAY;
180     char* host;
181     n=(unsigned long)-1;
182     host=argv[optind]+1;
183     {
184       int tmp;
185       tmp=str_chr(host,'/');
186       if (host[tmp]) {
187 	host[tmp]=0;
188 	if (!scan_ushort(host+tmp+1,&port)) usage();
189       }
190       tmp=str_chr(host,'%');
191       if (host[tmp]) {
192 	host[tmp]=0;
193 	scope_id=socket_getifidx(host+tmp+1);
194 	if (scope_id==0)
195 	  carp("warning: network interface \"",host+tmp+1,"\" not found.");
196       }
197     }
198 
199     {
200       stralloc a={0};
201       stralloc_copys(&a,host);
202       if (dns_ip6(&ips,&a)==-1)
203 	die(1,"could not find IP for \"",host,"\"!");
204     }
205     hostname=host;
206     request=krequest=0;
207     rlen=krlen=0;
208   } else {
209     char* host=argv[optind];
210     int colon;
211     int slash;
212     char* c;
213     mode=SAME;
214     if (byte_equal(host,7,"http://")) host+=7;
215     colon=str_chr(host,':');
216     slash=str_chr(host,'/');
217     if (host[0]=='[') {	/* ipv6 IP notation */
218       int tmp;
219       ++host;
220       --colon; --slash;
221       tmp=str_chr(host,']');
222       if (host[tmp]==']') host[tmp]=0;
223       if (host[tmp+1]==':') colon=tmp+1;
224       if (colon<tmp+1) colon=tmp+1+str_len(host+tmp+1);
225     }
226     if (colon<slash) {
227       host[colon]=0;
228       c=host+colon+1;
229       if (c[scan_ushort(c,&port)]!='/') usage();
230       *c=0;
231     }
232     host[colon]=0;
233     c=host+slash;
234     *c=0;
235     {
236       char* tmp=alloca(str_len(host)+1);
237       tmp[fmt_str(tmp,host)]=0;
238       host=tmp;
239     }
240     *c='/';
241     {
242       int tmp=str_chr(host,'%');
243       if (host[tmp]) {
244 	host[tmp]=0;
245 	scope_id=socket_getifidx(host+tmp+1);
246 	if (scope_id==0)
247 	  carp("warning: network interface \"",host+tmp+1,"\" not found.");
248       }
249     }
250 
251     {
252       stralloc a={0};
253       stralloc_copys(&a,host);
254       if (dns_ip6(&ips,&a)==-1)
255 	die(1,"could not find IP for \"",host,"\"!");
256     }
257 
258     request=alloca(1300+str_len(host)+3*str_len(c));
259     krequest=alloca(1300+str_len(host)+3*str_len(c));
260     {
261       int i,j;
262       i=fmt_str(request,"GET ");
263       i+=fmt_urlencoded(request+i,c,str_len(c));
264       i+=fmt_str(request+i," HTTP/1.0\r\nHost: ");
265       i+=fmt_str(request+i,host);
266       i+=fmt_str(request+i,":");
267       i+=fmt_ulong(request+i,port);
268       i+=fmt_str(request+i,"\r\nUser-Agent: bench/1.0\r\nConnection: ");
269       j=i;
270       i+=fmt_str(request+i,"close\r\n\r\n");
271       rlen=i; request[rlen]=0;
272       byte_copy(krequest,rlen,request);
273       i=j+fmt_str(krequest+j,"keep-alive\r\n\r\n");
274       krlen=i; krequest[krlen]=0;
275     }
276 
277     hostname=host;
278 
279   }
280 
281   fds=alloca(c*sizeof(*fds));
282   avail=alloca(c*sizeof(*avail));
283   expected=alloca(c*sizeof(*expected));
284   keepleft=alloca(c*sizeof(*keepleft));
285   requests=alloca(c*sizeof(size_t));
286   last.sec.x=23;
287   if (!k) K=1;
288   for (i=0; i<c; ++i) { fds[i]=-1; avail[i]=1; keepleft[i]=K; requests[i]=0; }
289 
290   taia_now(&first);
291 
292   for (done=0; done<n; ) {
293 
294     if (t) {
295       /* calculate timeout */
296       taia_now(&now);
297       if (now.sec.x != last.sec.x) {
298 	byte_copy(&last,sizeof(now),&now);
299 	byte_copy(&next,sizeof(now),&now);
300 	next.sec.x += t;
301 	while ((i=io_timeouted())!=-1) {
302 	  unsigned long j;
303 	  char numbuf[FMT_ULONG];
304 	  numbuf[fmt_ulong(numbuf,i)]=0;
305 	  carp("timeout on fd ",numbuf,"!");
306 	  j=(unsigned long)io_getcookie(i);
307 	  io_close(i);
308 	  avail[j]=1;
309 	  fds[j]=-1;
310 	}
311       }
312     }
313 
314     /* first, fill available connections */
315     for (i=0; i<c; ++i)
316       if (avail[i]==1 && fds[i]==-1) {
317 	fds[i]=make_connection(ips.s,port,scope_id,-1);
318 	if (fds[i]==-1) diesys(1,"socket/connect");
319 	avail[i]=2;
320 #ifdef HAVE_IO_FD_FLAGS
321 	if (io_fd_flags(fds[i],IO_FD_CANWRITE|IO_FD_NONBLOCK)==0) diesys(1,"io_fd");
322 #else
323 	if (io_fd_canwrite(fds[i])==0) diesys(1,"io_fd");
324 #endif
325 	io_setcookie(fds[i],(void*)i);
326 //	io_wantread(fds[i]);
327 	io_wantwrite(fds[i]);
328       }
329 
330     if (t)
331       io_waituntil(next);
332     else
333       io_wait();
334 
335     /* second, see if we can write on a connection */
336     while ((i=io_canwrite())!=-1) {
337       int j;
338       j=(unsigned long)io_getcookie(i);
339       if (avail[j]==2) {
340 	if (make_connection(ips.s,port,scope_id,i)==-1) {
341 	  ++errors;
342 	  if (v) write(1,"!",1);
343 	  io_close(i);
344 	  avail[j]=1;
345 	  fds[j]=-1;
346 	  continue;
347 	}
348       }
349       {
350 	char* towrite;
351 	int writelen;
352 	if (mode==REPLAY) {
353 	  static long lines;
354 	  char line[1024];
355 	  char req[2048];
356 	  int len;
357 	  int i;
358 	  char* c;
359 	  char* host;
360 	  int hlen;
361 	  if ((len=buffer_getline(buffer_0,line,sizeof(line)))) {
362 	    ++lines;
363 	    if (line[len]!='\n')
364 	      die(0,"line too long: ",line);
365 	    line[len]=0;
366 	    c=line;
367 	    if (str_start(line,"http://")) c+=7;
368 	    if (c[0]=='/') {
369 	      host=hostname;
370 	      hlen=strlen(hostname);
371 	    } else {
372 	      host=c;
373 	      c+=(hlen=str_chr(c,'/'));
374 	    }
375 	    if (!*c)
376 	      c="/";
377 
378 	    i=fmt_str(req,"GET ");
379 	    i+=fmt_urlencoded(req+i,c,str_len(c));
380 	    i+=fmt_str(req+i," HTTP/1.0\r\nHost: ");
381 	    byte_copy(req+i,hlen,host); i+=hlen;
382 	    i+=fmt_str(req+i,":");
383 	    i+=fmt_ulong(req+i,port);
384 	    if (cookies) {
385 	      int j;
386 	      i+=fmt_str(req+i,"\r\n");
387 	      j=nextcookie(req+i,sizeof(req)-i-100);
388 	      if (j!=-1) i+=j; else i-=2;
389 	    }
390 	    i+=fmt_str(req+i,"\r\nUser-Agent: bench/1.0\r\nConnection: ");
391 	    i+=fmt_str(req+i,keepleft[j]>1?"keep-alive\r\n\r\n":"close\r\n\r\n");
392 	    req[i]=0;
393 	    towrite=req;
394 	    writelen=i;
395 	  } else {
396 	    n=done;
397 	    break;
398 	  }
399 	} else {
400 	  if (keepleft[j]>1) {
401 	    towrite=krequest;
402 	    writelen=krlen;
403 	  } else {
404 	    towrite=request;
405 	    writelen=rlen;
406 	  }
407 	  if (cookies) {
408 	    int i=writelen-2;
409 	    int j=nextcookie(towrite+i,900);
410 	    if (j!=-1) i+=j;
411 	    i+=fmt_str(towrite+i,"\r\n\r\n");
412 	    writelen=i;
413 	  }
414 	}
415 	if (io_trywrite(i,towrite,writelen)!=writelen) {
416 	  ++errors;
417 	  if (v) write(1,"-",1);
418 	  io_close(i);
419 	  avail[j]=1;
420 	  fds[j]=-1;
421 	  continue;
422 	}
423       }
424       io_dontwantwrite(i);
425       io_wantread(i);
426       expected[j]=-1;
427       if (v) write(1,"+",1);
428     }
429 
430     /* third, see if we got served */
431     while ((i=io_canread())!=-1) {
432       char buf[8193];
433       int l,j;
434       buf[8192]=0;
435       j=(unsigned long)io_getcookie(i);
436       if ((l=io_tryread(i,buf,sizeof(buf)-1))<=0) {
437 	if (l==0) { /* EOF.  Mhh. */
438 	  if (expected[j]>0) {
439 	    ++errors;
440 	    if (v) write(1,"-",1);	/* so whine a little */
441 	  }
442 	  if (expected[j]==-2)
443 	    ++done;
444 	  io_close(i);
445 	  avail[j]=1;
446 	  fds[j]=-1;
447 	} else if (l==-3) {
448 	  ++errors;
449 	  if (v) write(1,"!",1);
450 //	  carpsys("read");
451 	}
452       } else {
453 	bytes+=l;
454 	if (v) write(1,".",1);
455 	/* read something */
456 	if (expected[j]==-1) {	/* expecting header */
457 	  int k;
458 	  /* OK, so this is a very simplistic header parser.  No
459 	   * buffering.  At all.  We expect the Content-Length header to
460 	   * come in one piece. */
461 	  if (l>10 && !memcmp(buf,"HTTP/1.",7)) {
462 	    if (buf[9]>='0' && buf[9]<='9') {
463 	      r[buf[9]-'0']++;
464 	      requests[j]++;
465 	    } else {
466 	      write(1,buf,15); write(1,"\n",1);
467 	    }
468 	  }
469 	  expected[j]=-2;
470 	  if (!done) {
471 	    for (k=0; k<l; ++k)
472 	      if (str_start(buf+k,"\nServer: ")) {
473 		char* tmp=buf+(k+=9);
474 		for (; k<l; ++k)
475 		  if (buf[k]=='\r') break;
476 		k=buf+k-tmp;
477 		if (k>sizeof(server)-1) k=sizeof(server)-1;
478 		byte_copy(server,k,tmp);
479 		server[k]=0;
480 		break;
481 	      }
482 	  }
483 	  for (k=0; k<l; ++k) {
484 	    if (str_start(buf+k,"\nContent-Length: ")) {
485 	      k+=17;
486 	      if (buf[k+scan_ulonglong(buf+k,(unsigned long long*)expected+j)] != '\r')
487 		die(1,"parse error in HTTP header!");
488 	    } else if (str_start(buf+k,"\r\n\r\n"))
489 	      break;
490 	  }
491 	  if (expected[j]>0) {
492 	    if (l-(k+4)>expected[j])
493 	      expected[j]=0;
494 	    else
495 	      expected[j]-=l-(k+4);
496 	  }
497 	} else if (expected[j]==-2) {
498 	  /* no content-length header, eat everything until EOF */
499 	} else {
500 	  if (l>expected[j]) {
501 	    carp("got more than expected!");
502 	    expected[j]=0;
503 	  } else
504 	    expected[j]-=l;
505 	}
506 	if (expected[j]==0) {
507 	  ++done;	/* one down! */
508 	  avail[j]=1;
509 //	  printf("fd %d: keepleft[%d]=%d\n",i,j,keepleft[j]);
510 	  if (keepleft[j]>1) {
511 	    --keepleft[j];
512 	    io_dontwantread(i);
513 	    io_wantwrite(i);
514 	    expected[j]=0;
515 	  } else {
516 	    keepleft[j]=K;
517 	    io_close(i);
518 	    fds[j]=-1;
519 	  }
520 	}
521       }
522     }
523   }
524 
525   taia_now(&now);
526   taia_sub(&now,&now,&first);
527   {
528     char a[FMT_ULONG];
529     char b[FMT_ULONG];
530     char C[FMT_ULONG];
531     char d[FMT_ULONG];
532     char e[FMT_ULONG];
533     char f[FMT_ULONG];
534     char g[FMT_ULONG];
535     char h[FMT_ULONG];
536     char i[FMT_ULONG];
537     char j[FMT_ULONG];
538     unsigned long long l;
539     a[fmt_ulong(a,now.sec.x)]=0;
540     b[fmt_ulong0(b,(now.nano%1000000000)/100000,4)]=0;
541     C[fmt_ulong(C,done)]=0;
542     d[fmt_ulonglong(d,errors)]=0;
543     e[fmt_ulonglong(e,bytes)]=0;
544 
545     /* let's say bytes = 10 MB, time = 1.2 sec.
546     * then we want 10*1024*1024/1.2 == 8 MB/sec */
547     l = (now.sec.x * 1024) + now.nano/976562;
548     if (l) {
549       int i;
550       l=bytes/l;
551       if (report)
552 	i=fmt_ulong(f,l);
553       else {
554 	i=fmt_humank(f,l*1024);
555 	i+=fmt_str(f+i,"iB/sec");
556       }
557       f[i]=0;
558     } else
559       strcpy(f,"n/a");
560 
561     l = (now.sec.x * 1000) + now.nano/1000000;
562     l = l ? (done*10000) / l : 0;
563     g[fmt_ulong(g,l/10)]=0;
564     h[fmt_ulong(h,c)]=0;
565     i[fmt_ulong(i,K)]=0;
566     j[fmt_ulong(j,kaputt)]=0;
567 
568     if (server[0]) msg("Server: ",server);
569     if (report) {
570       errmsg_iam(0);
571       msg("req\terr\tconcur\tkeep\tkbytes\tsec\ttput\tr/s\treset");
572       msg(C,"\t",d,"\t",h,"\t",i,"\t",e,"\t",a,".",b,"\t",f,"\t",g,"\t",j);
573     } else {
574       msg(C," requests, ",d," errors.");
575       msg(e," bytes in ",a,".",b," seconds.");
576       msg("Throughput: ",f);
577       msg("Requests per second: ",g);
578       msg("Connection refused/reset by peer: ",j);
579     }
580 
581     {
582       int i;
583       for (i=0; i<9; ++i) {
584 	a[fmt_ulong(a,r[i])]=0;
585 	b[0]=i+'0'; b[1]=0;
586 	msg(b,"xx: ",a);
587       }
588     }
589   }
590 
591   if (p) {
592     char a[FMT_ULONG];
593     char b[FMT_ULONG];
594     msg("successful requests per fd:\n");
595     for (i=0; i<c; ++i) {
596       a[fmt_ulong(a,i)]=0;
597       b[fmt_ulong(b,requests[i])]=0;
598       msg(a," ",b);
599     }
600   }
601 
602   return 0;
603 }
604