写一个mongodb压力测试程序

闲着无事写一个mongodb的压力测试程序，支持设置数据库的并发数，数据长短，每个连接写数据的数目。（用到mongodb官方的C驱动mongoc库）（没有考虑连接失败的情况，只是个简单的并发测试例子，当n指定过大后会连接失败）

功能如下：

vm6245:~/penjin/mongodb # ./test -h
mongo-bench
-i <ip_addr> host of the mongo server
-p <port>    post of the mongo server
-n <num>     number of parrel connections (max:10000)
-c <num>     count of date
-s <size>    data size of value (max:10000)

运行结果如：

vm6245:~/penjin/mongodb # ./a.out -i 10.6.2.171 -n 1 -c 100000
cost time 10000 3379297.000000(ms) 2959.195359(ops)
cost time 20000 6718649.000000(ms) 2976.788935(ops)
cost time 30000 10248220.000000(ms) 2927.337625(ops)
cost time 40000 13575299.000000(ms) 2946.528102(ops)
cost time 50000 17083459.000000(ms) 2926.807738(ops)
cost time 60000 20395650.000000(ms) 2941.803767(ops)
cost time 70000 24786119.000000(ms) 2824.161378(ops)
cost time 80000 28308409.000000(ms) 2826.015408(ops)
cost time 90000 31794577.000000(ms) 2830.671407(ops)
cost time 100000 35340010.000000(ms) 2829.653981(ops)
Write DB Cost time(microsecond):35340126.000000
cost time 10000 1942816.000000(ms) 5147.167822(ops)
cost time 20000 3981457.000000(ms) 5023.286701(ops)
cost time 30000 6066384.000000(ms) 4945.285363(ops)
cost time 40000 8378608.000000(ms) 4774.062708(ops)
cost time 50000 10560843.000000(ms) 4734.470534(ops)
cost time 60000 12526382.000000(ms) 4789.890648(ops)
cost time 70000 14427227.000000(ms) 4851.937243(ops)
cost time 80000 16435556.000000(ms) 4867.495812(ops)
cost time 90000 18462319.000000(ms) 4874.793898(ops)
cost time 100000 20361408.000000(ms) 4911.251717(ops)
Read DB Cost time(microsecond):20361475.000000

Parallel Connections: 1
Data count:100000
Data length:100

程序代码如下：（编译方法gcc test.c -o test -lmongoc -std=c99 -lpthread）

#include<stdio.h>
#include<stdlib.h>
#include<string.h>
#include<pthread.h>
#include<getopt.h>
#include<time.h>
#include<unistd.h>
#include<sys/time.h>
#include"mongo.h"

char *mongo_host = "127.0.0.1";
int mongo_port = 27017;
int paral_num = 10;
int data_count = 10000;
int data_length = 100;

#define MAX_SIZE 10000
#define TEST_DB_NAME "ciaos.test"

struct timeval start,current;
int count = 0;
pthread_mutex_t mutex = PTHREAD_MUTEX_INITIALIZER;
int child_thread_id = 1;

void* write_db()
{
	int i,j;
	int current_thread_id;
	//connect
	mongo conn[1];
	int status = mongo_connect(conn, mongo_host, mongo_port);
//  	mongo_replset_init( conn, "shard1" );
//  	mongo_replset_add_seed( conn, "10.6.2.15", 27017 );
//  	mongo_replset_add_seed( conn, "10.6.2.142", 27017 );
//  	mongo_replset_add_seed( conn, "10.6.2.154", 27017 );
//  	mongo_replset_add_seed( conn, "10.6.2.171", 27017 );
//  	mongo_replset_add_seed( conn, "10.6.2.174", 27017 );
//  	mongo_replset_add_seed( conn, "10.6.2.178", 27017 );
	
//	int status = mongo_replset_connect(conn);

	if(status != MONGO_OK){
		switch(conn->err){
		        case MONGO_CONN_SUCCESS:    printf( "connection succeeded\n" ); return 0;
		        case MONGO_CONN_NO_SOCKET:  printf( "no socket\n" ); return 0;
		        case MONGO_CONN_FAIL:       printf( "connection failed\n" ); return 0;
			case MONGO_CONN_NOT_MASTER: printf( "not master\n" ); return 0;
		}
	}
	mongo_write_concern write_concern[1];
	mongo_write_concern_init(write_concern);
	write_concern->w = 1;
	mongo_write_concern_finish(write_concern);

	//combine bson data
	char key[100],val[MAX_SIZE];
	
	pthread_mutex_lock(&mutex);
	current_thread_id = child_thread_id;
	child_thread_id ++;
	pthread_mutex_unlock(&mutex);

	for(i = 0;i < 9999 && i < data_length;i ++)
	{
		val[i] = 'x';
	}
	val[i] = '\0';

	bson b[1];
	for(i = 0;i < data_count;i ++)
	{
		sprintf(key,"childthread%dkey%d",current_thread_id,i);

		bson_init(b);
		bson_append_string(b,"_id",key);
		bson_append_string(b,"val",val);
		bson_finish(b);

		mongo_insert(conn, TEST_DB_NAME, b, write_concern);

		pthread_mutex_lock(&mutex);
		count ++;
		if(count % 10000 == 0){
			gettimeofday(¤t,NULL);
			double microsecs = 1000000*(current.tv_sec-start.tv_sec)+(current.tv_usec-start.tv_usec);
			printf("cost time %d %lf(ms) %lf(ops)\n",count,microsecs,count/(microsecs/1000000));
		}
		pthread_mutex_unlock(&mutex);
	}
	bson_destroy(b);
	//disconnect
	mongo_destroy(conn);

	return 0;
}

void* read_db()
{
	int i;
	int current_thread_id;

        mongo conn[1];
        int status = mongo_connect(conn, mongo_host, mongo_port);
//      mongo_replset_init( conn, "shard1" );
//      mongo_replset_add_seed( conn, "10.6.2.15", 27017 );
//      mongo_replset_add_seed( conn, "10.6.2.142", 27017 );
//      mongo_replset_add_seed( conn, "10.6.2.154", 27017 );
//      int status = mongo_replset_connect(conn);

        if(status != MONGO_OK){
                switch(conn->err){
                        case MONGO_CONN_SUCCESS:    printf( "connection succeeded\n" ); return 0;
                        case MONGO_CONN_NO_SOCKET:  printf( "no socket\n" ); return 0;
                        case MONGO_CONN_FAIL:       printf( "connection failed\n" ); return 0;
                        case MONGO_CONN_NOT_MASTER: printf( "not master\n" ); return 0;
                }
        }

        char key[100];

	pthread_mutex_lock(&mutex);
	current_thread_id = child_thread_id;
	child_thread_id ++;
	pthread_mutex_unlock(&mutex);

	bson q[1];
	mongo_cursor cursor[1];
	for(i = 0;i < data_count;i ++)
	{
		sprintf(key,"childthread%dkey%d",current_thread_id,i);

		bson_init(q);
		bson_append_string(q,"_id",key);
		bson_finish(q);

		mongo_cursor_init(cursor, conn, TEST_DB_NAME);
		mongo_cursor_set_query(cursor,q);

		while(mongo_cursor_next(cursor)==MONGO_OK){
			bson_iterator iterator[1];
			if (bson_find(iterator, mongo_cursor_bson(cursor), key)){
    			   // printf("%s=%s\n",key , bson_iterator_string(iterator));
    			}
		}

		pthread_mutex_lock(&mutex);
		count ++;
		if(count % 10000 == 0){
			gettimeofday(¤t,NULL);
			double microsecs = 1000000*(current.tv_sec-start.tv_sec)+(current.tv_usec-start.tv_usec);
			printf("cost time %d %lf(ms) %lf(ops)\n",count,microsecs,count/(microsecs/1000000));
		}
		pthread_mutex_unlock(&mutex);
	}

	bson_destroy(q);
	mongo_cursor_destroy(cursor);
	//disconnect
	mongo_destroy(conn);

	return 0;
}

void create_threads(void *action)
{
	int i;
	int ret = 0;
	pthread_t pid[MAX_SIZE];

	for(i = 0;i < paral_num;i ++){
		ret = pthread_create(pid+i, NULL, action, NULL);
		if(ret != 0){
			printf("create thread %d error\r\n",i+1);
			return;
		}
	}

	for(i = 0;i < paral_num;i ++){
		pthread_join(pid[i],NULL);
	}
}

void help()
{
	char *h = 
	"mongo-bench\n"
	"-i <ip_addr> host of the mongo server\n"
	"-p <port>    post of the mongo server\n"
	"-n <num>     number of parrel connections (max:10000)(default:10)\n"
	"-c <num>     count of date(default:10000)\n"
	"-s <size>    data size of value (max:10000)(default 100)\n";

	fprintf(stderr, h, strlen(h));
}

int main(int argc, char *argv[])
{
	int c;
	double microsecs;

	//amalyze arguments
	while((c = getopt(argc,argv,"i:p:n:c:s:h")) != -1){
		switch(c){
		case 'i':
			mongo_host = strdup(optarg);
			break;
		case 'p':
			mongo_port = atoi(optarg);
			break;
		case 'n':
			paral_num = atoi(optarg);
			break;
		case 'c':
			data_count = atoi(optarg);
			break;
		case 's':
			data_length = atoi(optarg);
			break;
		case 'h':
		default:
			help();
			return 1;
		}
	}

	//test write
	gettimeofday(&start,NULL);
	create_threads(write_db);
	gettimeofday(¤t,NULL);
	microsecs = 1000000*(current.tv_sec-start.tv_sec)+(current.tv_usec-start.tv_usec);
	printf("Write DB Cost time(microsecond):%lf\n",microsecs);

	//test read
	child_thread_id = 1;
	count = 0;
	gettimeofday(&start,NULL);
	create_threads(read_db);
	gettimeofday(¤t,NULL);
	microsecs = 1000000*(current.tv_sec-start.tv_sec)+(current.tv_usec-start.tv_usec);
	printf("Read DB Cost time(microsecond):%lf\n",microsecs);

	printf("\nParallel Connections: %d\n",paral_num);
	printf("Data count:%d\n",data_count);
	printf("Data length:%d\n",data_length);

	//drop test collections
	mongo conn[1];
	int status = mongo_connect(conn, mongo_host, mongo_port);
	if(status == MONGO_OK){
		mongo_cmd_drop_collection(conn, "ciaos", "test", NULL );
		mongo_destroy(conn);
	}
	return 0;
}

实现测试集群也很简单，test_db函数中将连接mongo单机数据库改为mongo数据库集群，代码如下:

        mongo_replset_init( conn, "shard1" );
        mongo_replset_add_seed( conn, "10.6.2.15", 27017 );
        mongo_replset_add_seed( conn, "10.6.2.142", 27017 );
        mongo_replset_add_seed( conn, "10.6.2.154", 27017 );
        int status = mongo_replset_connect(conn);