diff --git a/src/client.c b/src/client.c index 76f9d69..bccf8ca 100644 --- a/src/client.c +++ b/src/client.c @@ -289,10 +289,21 @@ static int get_chunk_server(TinyDFS *tdfs, Address *addrs, int num_addrs, ByteQu if (!tdfs->chunk_servers[i].used) continue; - if (have_insertection(addrs, num_addrs, tdfs->chunk_servers[i].addrs, tdfs->chunk_servers[i].num_addrs)) { - found = i; - break; - } + if (!have_insertection(addrs, num_addrs, tdfs->chunk_servers[i].addrs, tdfs->chunk_servers[i].num_addrs)) + continue; + + // It's possible that this chunk server's connection + // was dropped but we still didn't process the DISCONNECT + // event. + int conn_idx = tcp_index_from_tag(&tdfs->tcp, found); + if (conn_idx < 0) + continue; + + if (output) + *output = tcp_output_buffer(&tdfs->tcp, conn_idx); + + found = i; + break; } if (found == -1) { @@ -322,14 +333,6 @@ static int get_chunk_server(TinyDFS *tdfs, Address *addrs, int num_addrs, ByteQu request_queue_init(&tdfs->chunk_servers[found].reqs); tdfs->num_chunk_servers++; - - } else { - - int conn_idx = tcp_index_from_tag(&tdfs->tcp, found); - assert(conn_idx > -1); - - if (output) - *output = tcp_output_buffer(&tdfs->tcp, conn_idx); } return found; diff --git a/src/metadata_server.c b/src/metadata_server.c index 4496911..e3c5f32 100644 --- a/src/metadata_server.c +++ b/src/metadata_server.c @@ -880,7 +880,7 @@ int metadata_server_step(MetadataServer *state, void **contexts, struct pollfd * switch (events[i].type) { case EVENT_CONNECT: - tcp_set_tag(&state->tcp, conn_idx, CONNECTION_TAG_UNKNOWN); + tcp_set_tag(&state->tcp, conn_idx, CONNECTION_TAG_UNKNOWN, false); break; case EVENT_DISCONNECT: @@ -914,9 +914,9 @@ int metadata_server_step(MetadataServer *state, void **contexts, struct pollfd * if (is_chunk_server_message_type(msg_type)) { int chunk_server_idx = state->num_chunk_servers++; chunk_server_peer_init(&state->chunk_servers[chunk_server_idx]); - tcp_set_tag(&state->tcp, conn_idx, chunk_server_idx); + tcp_set_tag(&state->tcp, conn_idx, chunk_server_idx, true); } else { - tcp_set_tag(&state->tcp, conn_idx, CONNECTION_TAG_CLIENT); + tcp_set_tag(&state->tcp, conn_idx, CONNECTION_TAG_CLIENT, false); } } diff --git a/src/tcp.c b/src/tcp.c index 4752aff..04a9b9c 100644 --- a/src/tcp.c +++ b/src/tcp.c @@ -401,8 +401,14 @@ void tcp_close(TCP *tcp, int conn_idx) // if the output buffer is empty, the connection should be closed here. } -void tcp_set_tag(TCP *tcp, int conn_idx, int tag) +void tcp_set_tag(TCP *tcp, int conn_idx, int tag, bool unique) { + assert(tag != -1); + + if (unique) + for (int i = 0; i < tcp->num_conns; i++) + assert(tcp->conns[i].tag != tag); + tcp->conns[conn_idx].tag = tag; } diff --git a/src/tcp.h b/src/tcp.h index c27fb2d..015bc7d 100644 --- a/src/tcp.h +++ b/src/tcp.h @@ -62,7 +62,7 @@ int tcp_register_events(TCP *tcp, void **contexts, struct pollfd *polled); ByteQueue *tcp_output_buffer(TCP *tcp, int conn_idx); int tcp_connect(TCP *tcp, Address addr, int tag, ByteQueue **output); void tcp_close(TCP *tcp, int conn_idx); -void tcp_set_tag(TCP *tcp, int conn_idx, int tag); +void tcp_set_tag(TCP *tcp, int conn_idx, int tag, bool unique); int tcp_get_tag(TCP *tcp, int conn_idx); #endif // TCP_INCLUDED