[BENCH] Allow specifying request block size

4 files changed, 32 insertions, 19 deletions

diff --git a/src/bench/connection.c b/src/bench/connection.c
index 03ad9e5..ce9438a 100644
--- a/src/bench/connection.c
+++ b/src/bench/connection.c
@@ -41,6 +41,7 @@ bool connection_init_n_times(
 		const char *lowerImage,
 		const uint16_t rid,
 		int ntimes,
+		int blockSize,
 		BenchCounters* counters
 		) {
 	for (int run_i = 0; run_i < ntimes; ++run_i) {
@@ -95,20 +96,31 @@ bool connection_init_n_times(
 			} else if ( rid != 0 && rid != remoteRid ) {
 				counters->fails++;
 				logadd( LOG_ERROR, "rid mismatch" );
-			} else if ( !dnbd3_get_block( sock, run_i * 4096, 4096, 0, 0 ) ) {
+			} else if ( !dnbd3_get_block( sock, run_i * blockSize, blockSize, 0, 0 ) ) {
 				counters->fails++;
 				logadd( LOG_ERROR, "send: get block failed" );
 			} else if ( !dnbd3_get_reply( sock, &reply ) ) {
 				counters->fails++;
 				logadd( LOG_ERROR, "recv: get block header failed" );
-			} else if ( recv( sock, trash, sizeof(trash), MSG_WAITALL|MSG_NOSIGNAL ) != sizeof(trash) ) {
-				counters->fails++;
-				logadd( LOG_ERROR, "recv: get block payload failed" );
 			} else {
-				counters->success++;
-				close( sock );
-				sock = -1;
-				continue;
+				int rv, togo = blockSize;
+				do {
+					rv = recv( sock, trash, MIN( sizeof(trash), togo ), MSG_WAITALL|MSG_NOSIGNAL );
+					if ( rv == -1 && errno == EINTR )
+						continue;
+					if ( rv <= 0 )
+						break;
+					togo -= rv;
+				} while ( togo > 0 );
+				if ( togo != 0 ) {
+					counters->fails++;
+					logadd( LOG_ERROR, "recv: get block payload failed (remaining %d)", togo );
+				} else {
+					counters->success++;
+					close( sock );
+					sock = -1;
+					continue;
+				}
 			}
 			// Failed
 			if ( sock != -1 ) {
diff --git a/src/bench/connection.h b/src/bench/connection.h
index ff71e15..69207ff 100644
--- a/src/bench/connection.h
+++ b/src/bench/connection.h
@@ -19,7 +19,7 @@ typedef struct _dnbd3_async {
 } dnbd3_async_t;
 
 
-bool connection_init_n_times(const char *hosts, const char *image, const uint16_t rid, int ntimes, BenchCounters* counters);
+bool connection_init_n_times(const char *hosts, const char *image, const uint16_t rid, int ntimes, int blockSize, BenchCounters* counters);
 
 bool connection_init(const char *hosts, const char *image, const uint16_t rid);
 
diff --git a/src/bench/helper.h b/src/bench/helper.h
index 8342a79..e0c0262 100644
--- a/src/bench/helper.h
+++ b/src/bench/helper.h
@@ -29,6 +29,7 @@ typedef struct BenchThreadData {
 	char* server_address;
 	char * image_name;
 	int runs;
+	int bs;
 	int threadNumber;
 	bool closeSockets;
 } BenchThreadData;
diff --git a/src/bench/main.c b/src/bench/main.c
index c86af81..f8c55c3 100644
--- a/src/bench/main.c
+++ b/src/bench/main.c
@@ -17,10 +17,6 @@
 #define debugf(...) do { logadd( LOG_DEBUG1, __VA_ARGS__ ); } while (0)
 
 
-/* Debug/Benchmark variables */
-static bool useDebug = false;
-
-
 static void printUsage(char *argv0, int exitCode)
 {
 	printf( "Usage: %s [--debug] --host <serverAddress(es)> --image <imageName> [--rid revision]\n", argv0 );
@@ -30,17 +26,18 @@ static void printUsage(char *argv0, int exitCode)
 	printf( "   -r --rid        Revision to use (omit or pass 0 for latest)\n" );
 	printf( "   -n --runs       Number of connection attempts per thread\n" );
 	printf( "   -t --threads    number of threads\n" );
-	printf( "   -l --log        Write log to given location\n" );
+	printf( "   -b --blocksize  Size of blocks to request (def. 4096)\n" );
 	exit( exitCode );
 }
 
-static const char *optString = "h:i:n:t:HvVd";
+static const char *optString = "b:h:i:n:t:Hv";
 static const struct option longOpts[] = {
         { "host", required_argument, NULL, 'h' },
         { "image", required_argument, NULL, 'i' },
         { "nruns", optional_argument, NULL, 'n' },
         { "threads", required_argument, NULL, 't' },
-        { "help", required_argument, NULL, 'H' },
+        { "blocksize", required_argument, NULL, 'b' },
+        { "help", no_argument, NULL, 'H' },
         { "version", no_argument, NULL, 'v' },
         { 0, 0, 0, 0 }
 };
@@ -60,6 +57,7 @@ void* runBenchThread(void* t) {
 			data->image_name,
 			0,
 			data->runs,
+			data->bs,
 			data->counter);
 	printf("Thread #%d finished\n", data->threadNumber);
 	return NULL;
@@ -74,6 +72,7 @@ int main(int argc, char *argv[])
 	bool closeSockets = false;
 	int n_runs = 100;
 	int n_threads = 1;
+	int bs = 4096;
 
 	if ( argc <= 1 || strcmp( argv[1], "--help" ) == 0 || strcmp( argv[1], "--usage" ) == 0 ) {
 		printUsage( argv[0], 0 );
@@ -93,15 +92,15 @@ int main(int argc, char *argv[])
 		case 't':
 			n_threads = atoi(optarg);
 			break;
+		case 'b':
+			bs = atoi(optarg);
+			break;
 		case 'c':
 			closeSockets = true;
 			break;
 		case 'H':
 			printUsage( argv[0], 0 );
 			break;
-		case 'd':
-			useDebug = true;
-			break;
 		default:
 			printUsage( argv[0], EXIT_FAILURE );
 		}
@@ -123,6 +122,7 @@ int main(int argc, char *argv[])
 			server_address,
 			image_Name,
 			n_runs,
+			bs,
 			i,
 			closeSockets};
 		threadData[i] = tmp2;