1 files changed, 29 insertions, 15 deletions

diff --git a/lib9p/srv.c b/lib9p/srv.c
index 3d990aa..9bfe19e 100644
--- a/lib9p/srv.c
+++ b/lib9p/srv.c
@@ -21,7 +21,7 @@
 
 #define NAME  reqmap
 #define KEY_T uint32_t
-#define VAL_T struct lib9p_reqctx
+#define VAL_T struct lib9p_srv_reqctx
 #define CAP   CONFIG_9P_MAX_REQS
 #include "map.h"
 
@@ -50,6 +50,7 @@ struct lib9p_sess {
 	uint32_t                 rerror_overhead;
 	/* mutable */
 	bool                     initialized;
+	bool                     closing;
 	struct reqmap            reqs;
 	struct fidmap            fids;
 };
@@ -157,7 +158,6 @@ COROUTINE lib9p_srv_read_cr(void *_srv) {
 			.max_msg_size    = CONFIG_9P_MAX_MSG_SIZE,
 			.rerror_overhead = initial_rerror_overhead,
 			.initialized     = false,
-			.refcount        = 1,
 		};
 		for (;;) {
 			/* Read the message.  */
@@ -190,15 +190,15 @@ COROUTINE lib9p_srv_read_cr(void *_srv) {
 				goto close;
 
 			/* Handle the message... in another coroutine.  */
-			sess.refcount++;
 			cr_chan_send(&srv->reqch, &req);
 			cr_pause_and_yield(); /* wait for it to have copied req */
 		}
 	close:
-		netio_close(conn.fd, true, (--sess.refcount) == 0);
-		if (sess.refcount) {
+		netio_close(conn.fd, true, sess.reqs.len == 0);
+		if (sess.reqs.len) {
+			sess.closing = true;
 			cr_pause_and_yield();
-			assert(sess.refcount == 0);
+			assert(sess.reqs.len == 0);
 			netio_close(conn.fd, false, true);
 		}
 	}
@@ -244,7 +244,8 @@ COROUTINE lib9p_srv_write_cr(void *_srv) {
 
 		handle_Tmessage(&req);
 
-		if ((--req.parent_sess->refcount) == 0)
+		reqmap_del(&req.parent_sess->reqs, req.tag);
+		if (req.parent_sess->closing && !req.parent_sess->reqs.len)
 			cr_unpause(req.parent_sess->parent_conn->reader);
 	}
 
@@ -327,14 +328,16 @@ static bool handle_Tmessage(struct lib9p_req *req) {
 	}
 
  write:
-	if (lib9p_ctx_has_error(&req->ctx))
+	if (lib9p_ctx_has_error(&req->ctx)) {
 		respond_error(req);
-	else {
+		return true;
+	} else {
 		if (lib9p_marshal(&req->ctx, typ+1, req->tag, host_resp, req->net_bytes))
 			goto write;
 		cr_mutex_lock(&req->parent_sess->parent_conn->writelock);
 		netio_write(req->parent_sess->parent_conn->fd, req->net_bytes, decode_u32le(req->net_bytes));
 		cr_mutex_unlock(&req->parent_sess->parent_conn->writelock);
+		return false;
 	}
 }
 
@@ -364,13 +367,24 @@ static void handle_Tversion(struct lib9p_req *ctx, struct lib9p_msg_Tversion *re
 		return;
 	}
 
-
-	ctx->parent_sess--;
-	MAP_FOREACH(&ctx->parent_sess->fids, fid, fptr) {
-		fptr->vtable.free(subctx, fptr->impldata);
-		fidmap_del(&ctx->parent_sess->fids, fid);
+	struct lib9p_srv_reqctx subctx = {
+		.base  = &ctx->ctx,
+	};
+	if (ctx->parent_sess->reqs.len) {
+		ctx->parent_sess->closing = true;
+		// TODO: send flush events
+		cr_pause_and_yield();
+		assert(ctx->parent_sess->reqs.len == 0);
 	}
-}
+	if (ctx->parent_sess->fids.len) {
+		uint32_t fid;
+		struct lib9p_srv_file *fptr;
+		MAP_FOREACH(&ctx->parent_sess->fids, fid, fptr) {
+			fptr->vtable.free(&subctx, fptr->impldata);
+			fidmap_del(&ctx->parent_sess->fids, fid);
+		}
+	}
+	// TODO: replace session with a new one?
 
 #pragma GCC diagnostic push
 #pragma GCC diagnostic ignored "-Wdiscarded-qualifiers"