Rework when a lookup is considered complete

Keep track of outstanding requests sent as part of a lookup and rely on
those individual request timeouts to signal the end of the lookup as
a whole. This also makes sure we have the lookup around for the entire
duration of all downstream requests.

4 files changed, 43 insertions, 34 deletions

diff --git a/src/main.c b/src/main.c
index b2e0a76..0f81440 100644
--- a/src/main.c
+++ b/src/main.c
@@ -176,6 +176,7 @@ int main(int argc, char** argv) {
 	proto_begin(&dht, time(NULL), &message_cursor, outbuff+32);
 	flush_messages(dht.sfd, outbuff, message_cursor);
 
+	time_t lookup_refresh = 0;
 	// Init the lookup
 	dht.lookup.target = (struct nodeid){.inner={0x19b8a941, 0x38fa0191, 0x1403fac2, 0x581000ab, 0x19583cda}};
 	dht.lookup.state = OP_PENDING; // We want this to run at some point.
@@ -188,6 +189,9 @@ int main(int argc, char** argv) {
 
 		bool timedout = false;
 		time_t next = dht.wake;
+		if(lookup_refresh != 0 && difftime(lookup_refresh, next) < 0.0) {
+			next = lookup_refresh;
+		}
 
 		if(next != 0) {
 			time_t sleepfor = next - time(NULL);
@@ -238,10 +242,16 @@ int main(int argc, char** argv) {
 		rc = proto_run(&dht, buff, recv_len, (struct sockaddr_in*)&remote, remote_len, now, &message_cursor, outbuff+OUTBOX_SIZE);
 		flush_messages(dht.sfd, outbuff, message_cursor);
 
-		if(dht.lookup.state == OP_COMPLETED && difftime(dht.lookup.timeout, now) < 0.0) {
-			dbg("Restart LOOKUP");
-			// Restart the lookup periodically
-			dht.lookup.state = OP_PENDING;
+		if(dht.lookup.state == OP_COMPLETED) {
+			if(lookup_refresh == 0) {
+				dbg("Lookup completed");
+				lookup_refresh = now + 3600;
+			} else if (difftime(lookup_refresh, now) < 0.0) {
+				dbg("Restart LOOKUP");
+				lookup_refresh = 0;
+				// Restart the lookup periodically
+				dht.lookup.state = OP_PENDING;
+			}
 		}
 
 		save_config();
diff --git a/src/proto.c b/src/proto.c
index 0f5669a..e85d45e 100644
--- a/src/proto.c
+++ b/src/proto.c
@@ -173,6 +173,7 @@ PROCESS_REPONSE(getclient_response);
 PROCESS_TIMEOUT(getclient_timeout);
 
 PROCESS_REPONSE(lookup_response);
+PROCESS_TIMEOUT(lookup_timeout);
 
 // Number of nodeid bits
 #define IDBITS 160
@@ -224,11 +225,12 @@ int send_lookup(struct dht* dht, struct nodeid* target, time_t now, const struct
 	memcpy(&message->dest, dest_addr, dest_len);
 	message->dest_len = dest_len;
 
+	dht->lookup.outstanding++;
 	dht->requestdata[reqId].cont.lookup = &dht->lookup;
 
 	dht->requestdata[reqId].fun = &lookup_response;
 	dht->requestdata[reqId].timeout = now + PROTO_TMOUT;
-	dht->requestdata[reqId].timeout_fun = NULL;
+	dht->requestdata[reqId].timeout_fun = &lookup_timeout;
 	memcpy(&dht->requestdata[reqId].addr, dest_addr, dest_len);
 	dht->requestdata[reqId].addr_len = dest_len;
 
@@ -282,6 +284,17 @@ int send_ping(struct dht* dht, struct nodeid* expected, time_t now, bool node_is
 	return 0;
 }
 
+PROCESS_TIMEOUT(lookup_timeout) {
+	// @HACK: This really sucks. maybe we should just pass in the request id
+	size_t reqId = (typeof(dht->requestdata[0])*)((void*)cont - offsetof(typeof(dht->requestdata[0]), cont)) - dht->requestdata;
+
+	assert(cont->lookup->outstanding != 0);
+	dbg("Lookup request %ld times out", reqId);
+	cont->lookup->outstanding--;
+
+	return 0;
+}
+
 PROCESS_REPONSE(lookup_response) {
 	struct benc_node stream[256];
 	struct bcursor bcursor;
@@ -291,11 +304,8 @@ PROCESS_REPONSE(lookup_response) {
 		fatal("Response too short");
 	}
 
-	if(cont->lookup->state != OP_ACTIVE) {
-		// The lookup is not running so we have nowhere to dump the result.
-		// Just discard the packet
-		return PROTO_EDISC;
-	}
+	assert(cont->lookup->outstanding > 0);
+	assert(cont->lookup->state == OP_ACTIVE);
 
 	struct nodeid id;
 	uint8_t nodes_len;
@@ -448,14 +458,14 @@ PROCESS_REPONSE(lookup_response) {
 
 		int rc = send_lookup(dht, &cont->lookup->target, now, (struct sockaddr*)&dest, sizeof(struct sockaddr_in), msgbuff);
 		if(rc == PROTO_ENOREQ) {
-			return rc;
+			cont->lookup->outstanding--;
+			return PROTO_ENOREQ;
 		} else if(rc != 0) {
 			fatal("failed %d", rc);
 		}
-
-		cont->lookup->timeout = now + 120;
 	}
 
+	cont->lookup->outstanding--;
 	return 0;
 }
 
@@ -858,11 +868,6 @@ static void recalulate_waketime(struct dht *dht) {
 	if(dht->wake == 0 || (req_timeout != 0 && difftime(req_timeout, dht->wake) < 0.0)) {
 		dht->wake = req_timeout;
 	}
-
-	prom_gauge_set(wakeup_time, dht->lookup.timeout, (const char *[]){"lookup"});
-	if(dht->wake == 0 || (dht->lookup.timeout != 0 && difftime(dht->lookup.timeout, dht->wake) < 0)) {
-		dht->wake = dht->lookup.timeout;
-	}
 }
 
 int proto_run(struct dht* dht, char* buff, size_t recv_len, struct sockaddr_in* remote, socklen_t remote_len, time_t now, struct message** output, const struct message* const output_end) {
@@ -947,17 +952,16 @@ int proto_run(struct dht* dht, char* buff, size_t recv_len, struct sockaddr_in*
 					prom_counter_inc(outbox_overflow, NULL);
 					dht->pause = true;
 					recalulate_waketime(dht);
-					return 0;
+					break;
 				} else if(rc != 0) {
 					fatal("NOPE %d", rc);
 				}
 			}
 
-			dht->lookup.timeout = now + 120;
 			dht->lookup.state = OP_ACTIVE;
 			prom_counter_inc(lookup_count, NULL);
-		} else if(dht->lookup.state == OP_ACTIVE && dht->lookup.timeout <= now) {
-			dht->lookup.timeout = now + 3600;
+			dbg("Lookup started");
+		} else if(dht->lookup.state == OP_ACTIVE && dht->lookup.outstanding == 0) {
 			dht->lookup.state = OP_COMPLETED;
 		}
 
diff --git a/src/proto.h b/src/proto.h
index 8b65d3d..638c1ec 100644
--- a/src/proto.h
+++ b/src/proto.h
@@ -35,7 +35,7 @@ struct lookup {
 	struct nodeid closest[8];
 	struct addr closest_addr[8];
 
-	time_t timeout;
+	uint64_t outstanding;
 };
 
 void token_create(struct tokens* tokens, time_t now, struct addr* remote, char* token);
diff --git a/test/proto.c b/test/proto.c
index e0bb9e9..29563a3 100644
--- a/test/proto.c
+++ b/test/proto.c
@@ -735,7 +735,7 @@ void test_lookup_response() {
 		TEST_ASSERT_EQUAL_PTR(message_cursor, outbuff+1);
 		TEST_ASSERT_EQUAL(91, outbuff[0].payload_len);
 		TEST_ASSERT_EQUAL_CHAR_ARRAY("d1:ad2:id20:BBBBBBBBBBBBBBBBBBBB6:target20:aaaaaaaaaaaaaaaaaaaae1:q9:find_node1:t1:11:y1:qe", outbuff[0].payload, 91);
-		TEST_ASSERT_EQUAL(now+120, dht.lookup.timeout);
+		TEST_ASSERT_EQUAL(1, dht.lookup.outstanding);
 
 		memcpy(&remote, &outbuff[0].dest, sizeof(remote));
 	}
@@ -757,14 +757,14 @@ void test_lookup_response() {
 		TEST_ASSERT_EQUAL_CHAR_ARRAY("d1:ad2:id20:BBBBBBBBBBBBBBBBBBBB6:target20:aaaaaaaaaaaaaaaaaaaae1:q9:find_node1:t1:21:y1:qe", outbuff[0].payload, 91);
 		TEST_ASSERT_EQUAL(91, outbuff[1].payload_len);
 		TEST_ASSERT_EQUAL_CHAR_ARRAY("d1:ad2:id20:BBBBBBBBBBBBBBBBBBBB6:target20:aaaaaaaaaaaaaaaaaaaae1:q9:find_node1:t1:31:y1:qe", outbuff[1].payload, 91);
-		TEST_ASSERT_EQUAL(now+120, dht.lookup.timeout);
+		TEST_ASSERT_EQUAL(2, dht.lookup.outstanding); // Resolve 1, add 2
 
 		memcpy(&remote, &outbuff[0].dest, sizeof(remote));
 	}
 
 	// In actual use, we would have a lot more network/public api traffic here
 	// that would respond to some more pings. I don't want to write that code,
-	// so instead I'll just fill out some of the internal structured myself.
+	// so instead I'll just fill out some of the internal structures myself.
 	//
 	// What we are emulating is that a bunch of nodes responded before the one
 	// we just asserted above. Those nodes happened to be closer to the final
@@ -802,9 +802,7 @@ void test_lookup_response() {
 		// And we didn't fan out to the new node since it's already part of the current frontier
 		TEST_ASSERT_EQUAL_PTR(message_cursor, outbuff);
 
-		// Timeout isn't updated since we didn't send anything
-		// @FRAGILE this has to match the step size of now
-		TEST_ASSERT_EQUAL(now+120-1, dht.lookup.timeout);
+		TEST_ASSERT_EQUAL(1, dht.lookup.outstanding); // Resolve 1
 	}
 
 	// No other packets arrive for this lookup command and it should time out
@@ -816,8 +814,7 @@ void test_lookup_response() {
 		proto_run(&dht, NULL, 0, (struct sockaddr_in*)&remote, sizeof(remote), now, &message_cursor, outbuff+2);
 
 		TEST_ASSERT_EQUAL(OP_COMPLETED, dht.lookup.state);
-		// Timeout was set
-		TEST_ASSERT_EQUAL(now+3600, dht.lookup.timeout);
+		TEST_ASSERT_EQUAL(0, dht.lookup.outstanding);
 	}
 
 	now += 10;
@@ -839,9 +836,7 @@ void test_lookup_response() {
 		// And we didn't fan out to the new node since the lookup is done.
 		TEST_ASSERT_EQUAL_PTR(message_cursor, outbuff);
 
-		// Timeout shouldn't have been update either
-		// @FRAGILE this has to match the step size of now
-		TEST_ASSERT_EQUAL(now-10+3600, dht.lookup.timeout);
+		TEST_ASSERT_EQUAL(0, dht.lookup.outstanding);
 	}
 
 	proto_end(&dht);