added receive worker

4 files changed, 328 insertions, 183 deletions

diff --git a/src/kernel/core.c b/src/kernel/core.c
index 1b0aa46..f4b7204 100644
--- a/src/kernel/core.c
+++ b/src/kernel/core.c
@@ -117,7 +117,7 @@ again:
 		goto out;
 	}
 
-	ret = dnbd3_send_request(dev, sock, blk_mq_rq_from_pdu(cmd));
+	ret = dnbd3_send_request(sock, blk_mq_rq_from_pdu(cmd), cmd);
 	if (ret == -EAGAIN) {
 		dev_err_ratelimited(disk_to_dev(dev->disk), "request send failed, requeueing\n");
 		dnbd3_requeue_cmd(cmd);
@@ -338,11 +338,16 @@ int dnbd3_add_device(dnbd3_device *dev, int minor)
 	struct gendisk *disk;
 	struct request_queue *q;
 	int err = -ENOMEM;
+	int i;
 	printk(KERN_DEBUG "dnbd3: adding device %i\n", minor);
 
 	mutex_init(&dev->device_lock);
 	mutex_lock(&dev->device_lock);
 
+	for (i = 0; i < NUMBER_CONNECTIONS; i++) {
+		dev->socks[i].device = dev;
+	}
+
 	disk = alloc_disk(1);
 	if (!disk) {
 		printk(KERN_WARNING "dnbd3: alloc_disc failed device %i\n", minor);
@@ -428,7 +433,7 @@ static int __init dnbd3_init(void)
 	int i;
 	printk(KERN_DEBUG "dnbd3: starting kernel module\n");
 
-	dnbd3_wq = alloc_workqueue("kdnbd3", WQ_MEM_RECLAIM | WQ_FREEZABLE, 1);
+	dnbd3_wq = alloc_workqueue("kdnbd3", WQ_MEM_RECLAIM | WQ_FREEZABLE | WQ_UNBOUND, 0);
 
 	if (max_devs < 0) {
 		printk(KERN_ERR "dnbd3: max_devs must be >= 0\n");
diff --git a/src/kernel/dnbd3.h b/src/kernel/dnbd3.h
index 787c7f1..4e124e0 100644
--- a/src/kernel/dnbd3.h
+++ b/src/kernel/dnbd3.h
@@ -50,10 +50,12 @@ typedef struct dnbd3_sock {
 	struct request *pending;
 	struct dnbd3_server *server;
 	uint32_t heartbeat_count;
+	uint32_t cookie;
 	uint8_t panic, discover, panic_count;
-//	struct dnbd3_device *container_of;
+	struct dnbd3_device *device;
 	struct work_struct keepalive;
 	struct timer_list keepalive_timer;
+	struct work_struct receive;
 } dnbd3_sock;
 
 typedef struct dnbd3_device {
@@ -111,11 +113,11 @@ typedef struct dnbd3_device {
 typedef struct dnbd3_cmd {
 	struct dnbd3_device *dnbd3;
 	struct mutex lock;
-	int index;
-	int cookie;
+//	int index;
+	uint32_t cookie;
 	blk_status_t status;
 	unsigned long flags;
-	uint32_t cmd_cookie;
+//	uint32_t cmd_cookie;
 } dnbd3_cmd;
 
 
diff --git a/src/kernel/net.c b/src/kernel/net.c
index ef3988e..0721803 100644
--- a/src/kernel/net.c
+++ b/src/kernel/net.c
@@ -22,6 +22,7 @@
 
 
 #include <net/sock.h>
+#include <linux/wait.h>
 
 #include "dnbd3.h"
 #include "clientconfig.h"
@@ -46,6 +47,8 @@
         h.msg_flags = MSG_WAITALL | MSG_NOSIGNAL; \
 	} while (0)
 
+static DECLARE_WAIT_QUEUE_HEAD(send_wq);
+static uint64_t send_wq_handle;
 
 static int dnbd3_socket_connect(dnbd3_device *dev, dnbd3_server *server);
 static int dnbd3_socket_disconnect(dnbd3_device *dev, dnbd3_server *server, dnbd3_sock *sock);
@@ -70,23 +73,28 @@ static void print_server_list(struct dnbd3_device *dev)
 	}
 }
 
-int dnbd3_send_request(struct dnbd3_device *dev, struct dnbd3_sock *sock, struct request *req)
+static uint64_t to_handle(uint32_t arg0, uint32_t arg1) {
+	return ((uint64_t) arg0 << 32) | arg1;
+}
+
+static uint32_t arg0_from_handle(uint64_t handle) {
+	return (uint32_t)(handle >> 32);
+}
+
+static uint32_t arg1_from_handle(uint64_t handle) {
+	return (uint32_t) handle;
+}
+
+int dnbd3_send_request(struct dnbd3_sock *sock, struct request *req, struct dnbd3_cmd *cmd)
 {
 	dnbd3_request_t dnbd3_request;
-	dnbd3_reply_t dnbd3_reply;
 	struct msghdr msg;
 	struct kvec iov[2];
 	size_t iov_num = 1;
 	size_t send_len;
-	struct req_iterator iter;
-	struct bio_vec bvec_inst;
-	struct bio_vec *bvec = &bvec_inst;
-	sigset_t blocked, oldset;
-	void *kaddr;
-	int result, count, remaining;
-	uint16_t rid;
-	uint64_t reported_size;
-	char *name;
+	int result;
+	uint32_t tag;
+	uint64_t handle;
 	serialized_buffer_t payload_buffer;
 	sock->pending = req;
 	init_msghdr(msg);
@@ -106,12 +114,12 @@ int dnbd3_send_request(struct dnbd3_device *dev, struct dnbd3_sock *sock, struct
 		dnbd3_request.size = 0;
 		break;
 	case DNBD3_REQ_OP_CONNECT:
-		printk(KERN_DEBUG "dnbd3: request operation connect to %s\n", dev->imgname);
+		printk(KERN_DEBUG "dnbd3: request operation connect to %s\n", sock->device->imgname);
 		dnbd3_request.cmd = CMD_SELECT_IMAGE;
 		serializer_reset_write(&payload_buffer);
 		serializer_put_uint16(&payload_buffer, PROTOCOL_VERSION);
-		serializer_put_string(&payload_buffer, dev->imgname);
-		serializer_put_uint16(&payload_buffer, dev->rid);
+		serializer_put_string(&payload_buffer, sock->device->imgname);
+		serializer_put_uint16(&payload_buffer, sock->device->rid);
 		serializer_put_uint8(&payload_buffer, 0); // is_server = false
 		iov[1].iov_base = &payload_buffer;
 		dnbd3_request.size = iov[1].iov_len = serializer_get_written_length(&payload_buffer);
@@ -120,172 +128,295 @@ int dnbd3_send_request(struct dnbd3_device *dev, struct dnbd3_sock *sock, struct
 	default:
 		return -EIO;
 	}
+	sock->cookie++;
+	if (cmd != NULL) {
+		cmd->cookie = sock->cookie;
+		tag = blk_mq_unique_tag(req);
+		handle = ((uint64_t) tag << 32) | sock->cookie;
+	} else {
+		handle = sock->cookie;
+	}
+	memcpy(&dnbd3_request.handle, &handle, sizeof(handle));
+	printk(KERN_DEBUG "dnbd3: request handle is %llu\n", dnbd3_request.handle);
 
-	dnbd3_request.handle = (uint64_t)(uintptr_t)req; // Double cast to prevent warning on 32bit
+//	dnbd3_request.handle = (uint64_t)(uintptr_t)req; // Double cast to prevent warning on 32bit
 	fixup_request(dnbd3_request);
 	iov[0].iov_base = &dnbd3_request;
 	iov[0].iov_len = sizeof(dnbd3_request);
 	send_len = iov_num == 1 ? sizeof(dnbd3_request) : iov[0].iov_len + iov[1].iov_len;
-	if (kernel_sendmsg(sock->sock, &msg, iov, iov_num, send_len) != send_len) {
+	if ((result = kernel_sendmsg(sock->sock, &msg, iov, iov_num, send_len)) != send_len) {
 		printk(KERN_ERR "dnbd3: connection to server lost\n");
-		result = -EIO;
 		goto error;
 	}
 
-	// receive net reply
-	iov[0].iov_base = &dnbd3_reply;
-	iov[0].iov_len = sizeof(dnbd3_reply);
-	result = kernel_recvmsg(sock->sock, &msg, iov, 1, sizeof(dnbd3_reply), msg.msg_flags);
-	if (!result) {
-		printk(KERN_ERR "dnbd3: connection to server lost\n");
-		result = -EIO;
-		goto error;
+	sock->pending = NULL;
+	result = 0;
+error:
+	return result;
+}
+
 
+int dnbd3_send_request_blocking(struct dnbd3_sock *sock, int dnbd3_cmd)
+{
+	int result = 0;
+	uint64_t handle;
+	struct request *req = kmalloc(sizeof(struct request), GFP_ATOMIC );
+	printk(KERN_DEBUG "dnbd3: starting blocking request\n");
+	if (!req) {
+		printk(KERN_ERR "dnbd3: kmalloc failed\n");
+		goto error;
 	}
-	fixup_reply(dnbd3_reply);
 
-	// check error
-	if (dnbd3_reply.magic != dnbd3_packet_magic) {
-		printk(KERN_ERR "dnbd3: wrong magic packet\n");
-		result = -EIO;
+	switch (dnbd3_cmd) {
+	case CMD_KEEPALIVE:
+	case CMD_GET_SERVERS:
+		dnbd3_cmd_to_priv(req, dnbd3_cmd);
+		break;
+	case CMD_SELECT_IMAGE:
+		dnbd3_connect(req);
+		break;
+	default:
+		printk(KERN_WARNING "dnbd3: unsupported command for blocking %d\n", dnbd3_cmd);
+		result = -EINVAL;
 		goto error;
 	}
 
-	if (dnbd3_reply.cmd == 0) {
-		printk(KERN_ERR "dnbd3: command was 0\n");
-		result = -EIO;
+	mutex_lock(&sock->lock);
+	result = dnbd3_send_request(sock, req, NULL);
+	if (result) {
+		mutex_unlock(&sock->lock);
 		goto error;
 	}
+	send_wq_handle = 0;
+	handle = to_handle(sock->device->minor, dnbd3_cmd);
 
+	mutex_unlock(&sock->lock);
 
-	switch (dnbd3_reply.cmd) {
-	case CMD_GET_BLOCK:
-		rq_for_each_segment(bvec_inst, req, iter) {
-			siginitsetinv(&blocked, sigmask(SIGKILL));
-			sigprocmask(SIG_SETMASK, &blocked, &oldset);
+	printk(KERN_DEBUG "dnbd3: blocking request going to sleep wait for handle %llu\n", handle);
+	wait_event_interruptible(send_wq, handle == send_wq_handle);
+	printk(KERN_DEBUG "dnbd3: blocking request woke up with handle %llu\n", handle);
 
-			kaddr = kmap(bvec->bv_page) + bvec->bv_offset;
-			iov[0].iov_base = kaddr;
-			iov[0].iov_len = bvec->bv_len;
-			if (kernel_recvmsg(sock->sock, &msg, iov, 1, bvec->bv_len, msg.msg_flags) != bvec->bv_len) {
+
+error:
+	if (req) {
+		kfree(req);
+	}
+	return result;
+}
+
+static void dnbd3_receive_work(struct work_struct *work)
+{
+	struct dnbd3_sock *sock = container_of(work, struct dnbd3_sock, receive);
+	struct dnbd3_device *dev = sock->device;
+	struct request *req = NULL;
+	dnbd3_reply_t dnbd3_reply;
+	struct dnbd3_cmd *cmd;
+	struct msghdr msg;
+	struct kvec iov;
+	struct req_iterator iter;
+	struct bio_vec bvec_inst;
+	struct bio_vec *bvec = &bvec_inst;
+	sigset_t blocked, oldset;
+	void *kaddr;
+	uint32_t tag, cookie;
+	uint16_t hwq;
+	int result, count, remaining;
+	uint16_t rid;
+	uint64_t reported_size, handle;
+	char *name;
+	serialized_buffer_t payload_buffer;
+	init_msghdr(msg);
+
+	while(sock->sock) {
+		iov.iov_base = &dnbd3_reply;
+		iov.iov_len = sizeof(dnbd3_reply);
+		result = kernel_recvmsg(sock->sock, &msg, &iov, 1, sizeof(dnbd3_reply), msg.msg_flags);
+		if (!result) {
+			printk(KERN_ERR "dnbd3: connection to server lost\n");
+			result = -EIO;
+			goto error;
+
+		}
+		result = 0;
+		fixup_reply(dnbd3_reply);
+
+		// check error
+		if (dnbd3_reply.magic != dnbd3_packet_magic) {
+			printk(KERN_ERR "dnbd3: wrong magic packet\n");
+			result = -EIO;
+			goto error;
+		}
+
+		if (dnbd3_reply.cmd == 0) {
+			printk(KERN_ERR "dnbd3: command was 0\n");
+			result = -EIO;
+			goto error;
+		}
+
+
+		switch (dnbd3_reply.cmd) {
+		case CMD_GET_BLOCK:
+			printk(KERN_DEBUG "dnbd3: handle is %llu\n", dnbd3_reply.handle);
+			memcpy(&handle, &dnbd3_reply.handle, sizeof(handle));
+			cookie = (uint32_t) handle;
+			tag = (uint32_t)(handle >> 32);
+
+			hwq = blk_mq_unique_tag_to_hwq(tag);
+			if (hwq < dev->tag_set.nr_hw_queues)
+				req = blk_mq_tag_to_rq(dev->tag_set.tags[hwq], blk_mq_unique_tag_to_tag(tag));
+			if (!req || !blk_mq_request_started(req)) {
+				dev_err(disk_to_dev(dev->disk), "Unexpected reply (%d) %p\n", tag, req);
+				continue;
+			}
+			cmd = blk_mq_rq_to_pdu(req);
+
+			mutex_lock(&cmd->lock);
+			if (cmd->cookie != cookie) {
+				dev_err(disk_to_dev(dev->disk), "Double reply on req %p, cookie %u, handle cookie %u\n",
+					req, cmd->cookie, cookie);
+				mutex_unlock(&cmd->lock);
+				continue;
+			}
+
+
+			rq_for_each_segment(bvec_inst, req, iter) {
+				siginitsetinv(&blocked, sigmask(SIGKILL));
+				sigprocmask(SIG_SETMASK, &blocked, &oldset);
+
+				kaddr = kmap(bvec->bv_page) + bvec->bv_offset;
+				iov.iov_base = kaddr;
+				iov.iov_len = bvec->bv_len;
+				if (kernel_recvmsg(sock->sock, &msg, &iov, 1, bvec->bv_len, msg.msg_flags) != bvec->bv_len) {
+					kunmap(bvec->bv_page);
+					sigprocmask(SIG_SETMASK, &oldset, NULL );
+					printk(KERN_ERR "dnbd3: could not receive form net to block layer\n");
+					mutex_unlock(&cmd->lock);
+					continue;
+				}
 				kunmap(bvec->bv_page);
+
 				sigprocmask(SIG_SETMASK, &oldset, NULL );
-				printk(KERN_ERR "dnbd3: could not receive form net to block layer\n");
+			}
+			mutex_unlock(&cmd->lock);
+			blk_mq_end_request(req, 0);
+			break;
+		case CMD_GET_SERVERS:
+			printk(KERN_DEBUG "dnbd3: get servers received\n");
+			mutex_lock(&dev->device_lock);
+			if (!dev->use_server_provided_alts) {
+				remaining = dnbd3_reply.size;
+				goto consume_payload;
+			}
+			dev->new_servers_num = 0;
+			count = MIN(NUMBER_SERVERS, dnbd3_reply.size / sizeof(dnbd3_server_entry_t));
+
+			if (count != 0) {
+				iov.iov_base = dev->new_servers;
+				iov.iov_len = count * sizeof(dnbd3_server_entry_t);
+				if (kernel_recvmsg(sock->sock, &msg, &iov, 1, (count * sizeof(dnbd3_server_entry_t)), msg.msg_flags) != (count * sizeof(dnbd3_server_entry_t))) {
+					printk(KERN_ERR "dnbd3: failed to get servers\n");
+					mutex_unlock(&dev->device_lock);
+					goto error;
+				}
+				dev->new_servers_num = count;
+			}
+			// If there were more servers than accepted, remove the remaining data from the socket buffer
+			remaining = dnbd3_reply.size - (count * sizeof(dnbd3_server_entry_t));
+consume_payload:
+			while (remaining > 0) {
+				count = MIN(sizeof(dnbd3_reply), remaining); // Abuse the reply struct as the receive buffer
+				iov.iov_base = &dnbd3_reply;
+				iov.iov_len = count;
+				result = kernel_recvmsg(sock->sock, &msg, &iov, 1, count, msg.msg_flags);
+				if (result <= 0) {
+					printk(KERN_ERR "dnbd3: failed to receive payload from get servers\n");
+					mutex_unlock(&dev->device_lock);
+					goto error;
+				}
+				result = 0;
+			}
+			mutex_unlock(&dev->device_lock);
+			break;
+		case CMD_LATEST_RID:
+			if (dnbd3_reply.size != 2) {
+				printk(KERN_ERR "dnbd3: failed to get latest rid, wrong size\n");
 				goto error;
 			}
-			kunmap(bvec->bv_page);
-
-			sigprocmask(SIG_SETMASK, &oldset, NULL );
-		}
-		blk_mq_end_request(req, 0);
-		break;
-	case CMD_GET_SERVERS:
-		printk(KERN_DEBUG "dnbd3: get servers received\n");
-		mutex_lock(&dev->device_lock);
-		if (!dev->use_server_provided_alts) {
-			remaining = dnbd3_reply.size;
-			goto consume_payload;
-		}
-		dev->new_servers_num = 0;
-		count = MIN(NUMBER_SERVERS, dnbd3_reply.size / sizeof(dnbd3_server_entry_t));
-
-		if (count != 0) {
-			iov[0].iov_base = dev->new_servers;
-			iov[0].iov_len = count * sizeof(dnbd3_server_entry_t);
-			if (kernel_recvmsg(sock->sock, &msg, iov, 1, (count * sizeof(dnbd3_server_entry_t)), msg.msg_flags) != (count * sizeof(dnbd3_server_entry_t))) {
-				printk(KERN_ERR "dnbd3: failed to get servers\n");
-				mutex_unlock(&dev->device_lock);
+			printk(KERN_DEBUG "dnbd3: latest rid received\n");
+			iov.iov_base = &rid;
+			iov.iov_len = sizeof(rid);
+			if (kernel_recvmsg(sock->sock, &msg, &iov, 1, iov.iov_len, msg.msg_flags) <= 0) {
+				printk(KERN_ERR "dnbd3: failed to get latest rid\n");
 				goto error;
 			}
-			dev->new_servers_num = count;
-		}
-		// If there were more servers than accepted, remove the remaining data from the socket buffer
-		remaining = dnbd3_reply.size - (count * sizeof(dnbd3_server_entry_t));
-consume_payload:
-		while (remaining > 0) {
-			count = MIN(sizeof(dnbd3_reply), remaining); // Abuse the reply struct as the receive buffer
-			iov[0].iov_base = &dnbd3_reply;
-			iov[0].iov_len = count;
-			result = kernel_recvmsg(sock->sock, &msg, iov, 1, count, msg.msg_flags);
-			if (result <= 0) {
-				printk(KERN_ERR "dnbd3: failed to receive payload from get servers\n");
-				mutex_unlock(&dev->device_lock);
+			rid = net_order_16(rid);
+			printk("Latest rid of %s is %d (currently using %d)\n", dev->imgname, (int)rid, (int)dev->rid);
+			dev->update_available = (rid > dev->rid ? 1 : 0);
+			break;
+		case CMD_KEEPALIVE:
+			if (dnbd3_reply.size != 0) {
+				printk(KERN_ERR "dnbd3: got keep alive packet with payload\n");
+				goto error;
+			}
+			printk(KERN_DEBUG "dnbd3: keep alive received\n");
+			break;
+		case CMD_SELECT_IMAGE:
+			printk(KERN_DEBUG "dnbd3: select image received\n");
+			// receive reply payload
+			iov.iov_base = &payload_buffer;
+			iov.iov_len = dnbd3_reply.size;
+			if ((result = kernel_recvmsg(sock->sock, &msg, &iov, 1, dnbd3_reply.size, msg.msg_flags)) != dnbd3_reply.size) {
+				printk(KERN_ERR "dnbd3: could not read CMD_SELECT_IMAGE payload on handshake, size is %d and should be%d\n",
+						result, dnbd3_reply.size);
 				goto error;
 			}
 			result = 0;
-		}
-		mutex_unlock(&dev->device_lock);
-		break;
-	case CMD_LATEST_RID:
-		if (dnbd3_reply.size != 2) {
-			printk(KERN_ERR "dnbd3: failed to get latest rid, wrong size\n");
-			goto error;
-		}
-		printk(KERN_DEBUG "dnbd3: latest rid received\n");
-		iov[0].iov_base = &rid;
-		iov[0].iov_len = sizeof(rid);
-		if (kernel_recvmsg(sock->sock, &msg, iov, 1, iov[0].iov_len, msg.msg_flags) <= 0) {
-			printk(KERN_ERR "dnbd3: failed to get latest rid\n");
-			goto error;
-		}
-		rid = net_order_16(rid);
-		printk("Latest rid of %s is %d (currently using %d)\n", dev->imgname, (int)rid, (int)dev->rid);
-		dev->update_available = (rid > dev->rid ? 1 : 0);
-		break;
-	case CMD_KEEPALIVE:
-		if (dnbd3_reply.size != 0) {
-			printk(KERN_ERR "dnbd3: got keep alive packet with payload\n");
-			goto error;
-		}
-		printk(KERN_DEBUG "dnbd3: keep alive received\n");
-		break;
-	case CMD_SELECT_IMAGE:
-		printk(KERN_DEBUG "dnbd3: select image received\n");
-		// receive reply payload
-		iov[0].iov_base = &payload_buffer;
-		iov[0].iov_len = dnbd3_reply.size;
-		if (kernel_recvmsg(sock->sock, &msg, iov, 1, dnbd3_reply.size, msg.msg_flags) != dnbd3_reply.size) {
-			printk(KERN_ERR "dnbd3: could not read CMD_SELECT_IMAGE payload on handshake\n");
-			goto error;
-		}
-
-		// handle/check reply payload
-		serializer_reset_read(&payload_buffer, dnbd3_reply.size);
-		sock->server->protocol_version = serializer_get_uint16(&payload_buffer);
-		if (sock->server->protocol_version < MIN_SUPPORTED_SERVER) {
-			printk(KERN_ERR "dnbd3: server version is lower than min supported version\n");
-			goto error;
-		}
 
-		name = serializer_get_string(&payload_buffer);
-		rid = serializer_get_uint16(&payload_buffer);
-		if (dev->rid != rid && strcmp(name, dev->imgname) != 0) {
-			printk(KERN_ERR "dnbd3: server offers image '%s', requested '%s'\n", name, dev->imgname);
-			goto error;
-		}
+			// handle/check reply payload
+			serializer_reset_read(&payload_buffer, dnbd3_reply.size);
+			sock->server->protocol_version = serializer_get_uint16(&payload_buffer);
+			if (sock->server->protocol_version < MIN_SUPPORTED_SERVER) {
+				printk(KERN_ERR "dnbd3: server version is lower than min supported version\n");
+				goto error;
+			}
 
-		reported_size = serializer_get_uint64(&payload_buffer);
-		if (!dev->reported_size) {
-			if (reported_size < 4096) {
-				printk(KERN_ERR "dnbd3: reported size by server is < 4096\n");
+			name = serializer_get_string(&payload_buffer);
+			rid = serializer_get_uint16(&payload_buffer);
+			if (dev->rid != rid && strcmp(name, dev->imgname) != 0) {
+				printk(KERN_ERR "dnbd3: server offers image '%s', requested '%s'\n", name, dev->imgname);
 				goto error;
 			}
-			dev->reported_size = reported_size;
-			set_capacity(dev->disk, dev->reported_size >> 9); /* 512 Byte blocks */
-		} else if (dev->reported_size != reported_size) {
-			printk(KERN_ERR "dnbd3: reported size by server is %llu but should be %llu\n", reported_size, dev->reported_size);
-		}
 
-		break;
-	default:
-		printk("ERROR: Unknown command (Receive)\n");
-		break;
+			reported_size = serializer_get_uint64(&payload_buffer);
+			if (!dev->reported_size) {
+				if (reported_size < 4096) {
+					printk(KERN_ERR "dnbd3: reported size by server is < 4096\n");
+					goto error;
+				}
+				dev->reported_size = reported_size;
+				set_capacity(dev->disk, dev->reported_size >> 9); /* 512 Byte blocks */
+			} else if (dev->reported_size != reported_size) {
+				printk(KERN_ERR "dnbd3: reported size by server is %llu but should be %llu\n", reported_size, dev->reported_size);
+			}
 
-	}
-	sock->pending = NULL;
-	result = 0;
+			break;
+		default:
+			printk(KERN_WARNING "dnbd3: Unknown command (Receive)\n");
+			break;
+		}
 error:
-	return result;
+		handle = to_handle(dev->minor, dnbd3_reply.cmd);
+		printk(KERN_DEBUG "dnbd3: try to wake up queue with handle %llu\n", handle);
+		send_wq_handle = handle;
+		wake_up_interruptible(&send_wq);
+		if (result) {
+			printk(KERN_DEBUG "dnbd3: receive error happened %d\n", result);
+			break; //TODO for now need to handle errors
+		}
+		printk(KERN_DEBUG "dnbd3: receive completed, waiting for next receive\n");
+	}
+	printk(KERN_DEBUG "dnbd3: receive work queue is stopped\n");
+	dnbd3_socket_disconnect(dev, sock->server, sock);
 }
 
 
@@ -300,20 +431,19 @@ void dnbd3_keepalive(struct timer_list *arg)
 static void keepalive(struct work_struct *work)
 {
 	struct dnbd3_sock *sock = container_of(work, struct dnbd3_sock, keepalive);
-	struct request *req;
+//	struct request *req;
 	printk(KERN_DEBUG "dnbd3: starting keepalive worker\n");
-	mutex_lock(&sock->lock);
-	req = kmalloc(sizeof(struct request), GFP_ATOMIC );
+//	mutex_lock(&sock->lock);
+//	req = kmalloc(sizeof(struct request), GFP_ATOMIC );
 	// send keepalive
-	if (req) {
-		dnbd3_cmd_to_priv(req, CMD_KEEPALIVE);
-		dnbd3_send_request(NULL, sock, req); // we do not need the device for keepalive
-		kfree(req);
-	} else {
-		printk(KERN_WARNING "dnbd3: could not create keepalive request\n");
-	}
+//	if (req) {
+		dnbd3_send_request_blocking(sock, CMD_KEEPALIVE);
+////		kfree(req);
+//	} else {
+//		printk(KERN_WARNING "dnbd3: could not create keepalive request\n");
+//	}
 	++sock->heartbeat_count;
-	mutex_unlock(&sock->lock);
+//	mutex_unlock(&sock->lock);
 }
 
 void dnbd3_discovery(struct timer_list *arg)
@@ -328,22 +458,24 @@ static void discovery(struct work_struct *work)
 {
 	struct dnbd3_device *dev = container_of(work, struct dnbd3_device, discovery);
 	dnbd3_sock *sock = &dev->socks[0]; // we use the first sock for discovery
-	struct request *req;
+//	struct request *req;
 	int i, j;
 	struct dnbd3_server *existing_server, *free_server, *failed_server;
 	dnbd3_server_entry_t *new_server;
 	printk(KERN_DEBUG "dnbd3: starting discovery worker\n");
-	mutex_lock(&sock->lock);
-	req = kmalloc(sizeof(struct request), GFP_ATOMIC );
-	// send keepalive
-	if (req) {
-		dnbd3_cmd_to_priv(req, CMD_GET_SERVERS);
-		dnbd3_send_request(dev, sock, req);
-		kfree(req);
-	} else {
-		printk(KERN_WARNING "dnbd3: could not create get servers request\n");
-	}
-	mutex_unlock(&sock->lock);
+//	mutex_lock(&sock->lock);
+//	req = kmalloc(sizeof(struct request), GFP_ATOMIC );
+//	// send keepalive
+//	if (req) {
+//		dnbd3_cmd_to_priv(req, CMD_GET_SERVERS);
+		dnbd3_send_request_blocking(sock, CMD_GET_SERVERS);
+//		kfree(req);
+//	} else {
+//		printk(KERN_WARNING "dnbd3: could not create get servers request\n");
+//	}
+//	mutex_unlock(&sock->lock);
+
+	//TODO wait until something is received
 
 	printk(KERN_DEBUG "dnbd3: new server num is %d\n", dev->new_servers_num);
 	if (dev->new_servers_num) {
@@ -488,21 +620,27 @@ static int dnbd3_socket_connect(dnbd3_device *dev, dnbd3_server *server)
 
 	__dnbd3_socket_connect(server, sock);
 
-	mutex_lock(&sock->lock);
-	req = kmalloc(sizeof(*req), GFP_ATOMIC );
-	if (!req) {
-		printk(KERN_ERR "dnbd3: kmalloc failed\n");
-		goto error;
-	}
+//	mutex_lock(&sock->lock);
+//	req = kmalloc(sizeof(*req), GFP_ATOMIC );
+//	if (!req) {
+//		printk(KERN_ERR "dnbd3: kmalloc failed\n");
+//		goto error;
+//	}
 
+	// start the receiver
+	INIT_WORK(&sock->receive, dnbd3_receive_work);
+	queue_work(dnbd3_wq, &sock->receive);
 
-	dnbd3_connect(req);
-	result = dnbd3_send_request(dev, sock, req);
+//	dnbd3_connect(req);
+	result = dnbd3_send_request_blocking(sock, CMD_SELECT_IMAGE);
 	if (result) {
 		printk(KERN_ERR "dnbd3: connection to image %s failed\n", dev->imgname);
 		goto error;
 
 	}
+//	mutex_unlock(&sock->lock);
+
+	//TODO wait until connected
 
 	printk(KERN_DEBUG "dnbd3: connected to image %s, filesize %llu\n", dev->imgname, dev->reported_size);
 
@@ -515,7 +653,7 @@ static int dnbd3_socket_connect(dnbd3_device *dev, dnbd3_server *server)
 
 	mutex_unlock(&sock->lock);
 
-	kfree(req);
+//	kfree(req);
 	return 0;
 error:
 	if (sock->sock) {
diff --git a/src/kernel/net.h b/src/kernel/net.h
index 6cd0ffa..435e39c 100644
--- a/src/kernel/net.h
+++ b/src/kernel/net.h
@@ -24,7 +24,7 @@
 #define NET_H_
 
 
-int dnbd3_send_request(struct dnbd3_device *dev, struct dnbd3_sock *sock, struct request *req);
+int dnbd3_send_request(struct dnbd3_sock *sock, struct request *req, struct dnbd3_cmd *cmd);
 
 
 int dnbd3_net_connect(struct dnbd3_device *dev);