diff --git a/src/chunk_server.c b/src/chunk_server.c
index cd740ae..e844dc0 100644
--- a/src/chunk_server.c
+++ b/src/chunk_server.c
@@ -764,7 +764,7 @@ process_client_message(ChunkServer *state, int conn_idx, uint16_t type, ByteView
     return -1;
 }
 
-int chunk_server_init(ChunkServer *state, int argc, char **argv, void **contexts, struct pollfd *polled)
+int chunk_server_init(ChunkServer *state, int argc, char **argv, void **contexts, struct pollfd *polled, int *timeout)
 {
     (void) argc;
     (void) argv;
@@ -805,6 +805,7 @@ int chunk_server_init(ChunkServer *state, int argc, char **argv, void **contexts
 
     state->metadata_server_disconnect_time = 0;
 
+    *timeout = -1;  // No timeout needed for chunk server initially
     return tcp_register_events(&state->tcp, contexts, polled);
 }
 
@@ -816,7 +817,7 @@ int chunk_server_free(ChunkServer *state)
     return 0;
 }
 
-int chunk_server_step(ChunkServer *state, void **contexts, struct pollfd *polled, int num_polled)
+int chunk_server_step(ChunkServer *state, void **contexts, struct pollfd *polled, int num_polled, int *timeout)
 {
     Event events[MAX_CONNS+1];
     int num_events = tcp_translate_events(&state->tcp, events, contexts, polled, num_polled);
@@ -937,5 +938,6 @@ int chunk_server_step(ChunkServer *state, void **contexts, struct pollfd *polled
         }
     }
 
+    *timeout = -1;  // No timeout needed for chunk server
     return tcp_register_events(&state->tcp, contexts, polled);
 }
diff --git a/src/chunk_server.h b/src/chunk_server.h
index 0125f5e..6fd29cd 100644
--- a/src/chunk_server.h
+++ b/src/chunk_server.h
@@ -35,8 +35,8 @@ typedef struct {
     PendingDownloadList pending_download_list;
 } ChunkServer;
 
-int chunk_server_init(ChunkServer *state, int argc, char **argv, void **contexts, struct pollfd *polled);
+int chunk_server_init(ChunkServer *state, int argc, char **argv, void **contexts, struct pollfd *polled, int *timeout);
 int chunk_server_free(ChunkServer *state);
-int chunk_server_step(ChunkServer *state, void **contexts, struct pollfd *polled, int num_polled);
+int chunk_server_step(ChunkServer *state, void **contexts, struct pollfd *polled, int num_polled, int *timeout);
 
 #endif // CHUNK_SERVER_INCLUDED
diff --git a/src/metadata_server.c b/src/metadata_server.c
index 902641b..0fcd75d 100644
--- a/src/metadata_server.c
+++ b/src/metadata_server.c
@@ -750,7 +750,7 @@ static bool is_chunk_server_message_type(uint16_t type)
     return false;
 }
 
-int metadata_server_init(MetadataServer *state, int argc, char **argv, void **contexts, struct pollfd *polled)
+int metadata_server_init(MetadataServer *state, int argc, char **argv, void **contexts, struct pollfd *polled, int *timeout)
 {
     (void) argc;
     (void) argv;
@@ -778,6 +778,7 @@ int metadata_server_init(MetadataServer *state, int argc, char **argv, void **co
         return -1;
     }
 
+    *timeout = -1;  // No timeout needed for metadata server
     return tcp_register_events(&state->tcp, contexts, polled);
 }
 
@@ -788,7 +789,7 @@ int metadata_server_free(MetadataServer *state)
     return 0;
 }
 
-int metadata_server_step(MetadataServer *state, void **contexts, struct pollfd *polled, int num_polled)
+int metadata_server_step(MetadataServer *state, void **contexts, struct pollfd *polled, int num_polled, int *timeout)
 {
     Event events[MAX_CONNS+1];
     int num_events = tcp_translate_events(&state->tcp, events, contexts, polled, num_polled);
@@ -852,5 +853,6 @@ int metadata_server_step(MetadataServer *state, void **contexts, struct pollfd *
         }
     }
 
+    *timeout = -1;  // No timeout needed for metadata server
     return tcp_register_events(&state->tcp, contexts, polled);
 }
diff --git a/src/metadata_server.h b/src/metadata_server.h
index 05344ef..48d6943 100644
--- a/src/metadata_server.h
+++ b/src/metadata_server.h
@@ -48,8 +48,8 @@ typedef struct {
 
 } MetadataServer;
 
-int metadata_server_init(MetadataServer *state, int argc, char **argv, void **contexts, struct pollfd *polled);
+int metadata_server_init(MetadataServer *state, int argc, char **argv, void **contexts, struct pollfd *polled, int *timeout);
 int metadata_server_free(MetadataServer *state);
-int metadata_server_step(MetadataServer *state, void **contexts, struct pollfd *polled, int num_polled);
+int metadata_server_step(MetadataServer *state, void **contexts, struct pollfd *polled, int num_polled, int *timeout);
 
 #endif // METADATA_SERVER_INCLUDED
diff --git a/src/system.c b/src/system.c
index 0b06033..6a88aec 100644
--- a/src/system.c
+++ b/src/system.c
@@ -323,26 +323,19 @@ int spawn_simulated_process(char *args)
     void *contexts[MAX_CONNS+1];
     struct pollfd polled[MAX_CONNS+1];
     int num_polled;
+    int timeout = -1;
 
     current_process = process;
 
     switch (process->type) {
         case PROCESS_TYPE_METADATA_SERVER:
-            num_polled = metadata_server_init(&process->metadata_server, argc, argv, contexts, polled);
+            num_polled = metadata_server_init(&process->metadata_server, argc, argv, contexts, polled, &timeout);
             break;
         case PROCESS_TYPE_CHUNK_SERVER:
-            num_polled = chunk_server_init(&process->chunk_server, argc, argv, contexts, polled);
+            num_polled = chunk_server_init(&process->chunk_server, argc, argv, contexts, polled, &timeout);
             break;
         case PROCESS_TYPE_CLIENT:
-        {
-            int timeout = -1;
             num_polled = simulation_client_init(&process->simulation_client, argc, argv, contexts, polled, &timeout);
-            if (timeout < 0) {
-                process->wakeup_time = INVALID_TIME;
-            } else {
-                process->wakeup_time = current_time + timeout;
-            }
-        }
             break;
         default:
             num_polled = -1;
@@ -354,6 +347,12 @@ int spawn_simulated_process(char *args)
         // TODO
     }
 
+    if (timeout < 0) {
+        process->wakeup_time = INVALID_TIME;
+    } else {
+        process->wakeup_time = current_time + timeout;
+    }
+
     process_poll_array(process, contexts, polled, num_polled);
 
     processes[num_processes++] = process;
@@ -546,23 +545,16 @@ void update_simulation(void)
             }
         }
 
+        int timeout = -1;
         switch (current_process->type) {
             case PROCESS_TYPE_METADATA_SERVER:
-                num_polled = metadata_server_step(&current_process->metadata_server, contexts, polled, num_polled);
+                num_polled = metadata_server_step(&current_process->metadata_server, contexts, polled, num_polled, &timeout);
                 break;
             case PROCESS_TYPE_CHUNK_SERVER:
-                num_polled = chunk_server_step(&current_process->chunk_server, contexts, polled, num_polled);
+                num_polled = chunk_server_step(&current_process->chunk_server, contexts, polled, num_polled, &timeout);
                 break;
             case PROCESS_TYPE_CLIENT:
-            {
-                int timeout = -1;
                 num_polled = simulation_client_step(&current_process->simulation_client, contexts, polled, num_polled, &timeout);
-                if (timeout < 0) {
-                    current_process->wakeup_time = INVALID_TIME;
-                } else {
-                    current_process->wakeup_time = current_time + timeout;
-                }
-            }
                 break;
         }
 
@@ -570,6 +562,12 @@ void update_simulation(void)
             // TODO
         }
 
+        if (timeout < 0) {
+            current_process->wakeup_time = INVALID_TIME;
+        } else {
+            current_process->wakeup_time = current_time + timeout;
+        }
+
         process_poll_array(current_process, contexts, polled, num_polled);
 
         current_process = NULL;