@joystick.js/db-canary 0.0.0-canary.2209

package/src/server/cluster/worker.js

@@ -0,0 +1,537 @@
+/**
+ * @fileoverview Cluster worker process for JoystickDB server.
+ * Handles TCP connections, client authentication, and routes operations to master process.
+ * Manages socket connections, message parsing, and graceful shutdown in worker processes.
+ */
+
+import net from 'net';
+import op_types from '../lib/op_types.js';
+import { send_success, send_error, send_message } from '../lib/send_response.js';
+import { shutdown_write_queue } from '../lib/write_queue.js';
+import { create_message_parser, encode_message } from '../lib/tcp_protocol.js';
+import create_logger from '../lib/logger.js';
+import { initialize_database, cleanup_database } from '../lib/query_engine.js';
+import { handle_admin_operation, handle_ping_operation } from '../lib/operation_dispatcher.js';
+
+/**
+ * Cluster worker process that handles TCP connections and routes operations to master.
+ * Manages client connections, authentication, and message routing in a clustered environment.
+ */
+class ClusterWorker {
+  /**
+   * Creates a new cluster worker instance.
+   * Initializes connection management, authentication tracking, and sets up process event handlers.
+   */
+  constructor() {
+    this.server = null;
+    this.connections = new Map();
+    this.connection_count = 0;
+    this.settings = null;
+    this.port = null;
+    this.write_id_counter = 0;
+    this.pending_writes = new Map();
+    this.authenticated_clients = new Set();
+    this.heartbeat_interval = null;
+    
+    const { create_context_logger } = create_logger('worker');
+    this.log = create_context_logger({ worker_pid: process.pid });
+    
+    this.setup_worker();
+  }
+  
+  setup_worker() {
+    process.on('message', (message) => {
+      this.handle_master_message(message);
+    });
+    
+    process.on('SIGTERM', () => {
+      this.shutdown();
+    });
+    
+    process.on('SIGINT', () => {
+      this.shutdown();
+    });
+    
+    this.send_heartbeat();
+    
+    this.heartbeat_interval = setInterval(() => {
+      this.send_heartbeat();
+    }, 5000);
+    
+    // NOTE: Signal to master that worker is ready to receive config.
+    if (process.connected) {
+      process.send({ type: 'worker_ready' });
+    }
+  }
+  
+  handle_master_message(message) {
+    switch (message.type) {
+      case 'config':
+        this.handle_config(message);
+        break;
+        
+      case 'write_response':
+        this.handle_write_response(message);
+        break;
+        
+      case 'auth_response':
+        this.handle_auth_response(message);
+        break;
+        
+      case 'setup_response':
+        this.handle_setup_response(message);
+        break;
+        
+      case 'write_notification':
+        this.handle_write_notification(message);
+        break;
+        
+      case 'shutdown':
+        this.shutdown();
+        break;
+        
+      default:
+        this.log.warn('Unknown message type received from master', { 
+          message_type: message.type 
+        });
+    }
+  }
+  
+  handle_config(message) {
+    const incoming_master_id = message.data.master_id;
+    
+    // NOTE: Only handle config from the first master to avoid conflicts between multiple test masters.
+    if (this.master_id && this.master_id !== incoming_master_id) {
+      this.log.info('Worker already configured by different master, ignoring config message', { 
+        current_master_id: this.master_id,
+        incoming_master_id: incoming_master_id,
+        current_port: this.port, 
+        new_port: message.data.port
+      });
+      return;
+    }
+    
+    // NOTE: If already configured by the same master, ignore duplicate config.
+    if (this.port !== null && this.master_id === incoming_master_id) {
+      this.log.info('Worker already configured by same master, ignoring duplicate config message', { 
+        master_id: incoming_master_id,
+        current_port: this.port, 
+        new_port: message.data.port
+      });
+      return;
+    }
+    
+    this.log.info('Received config message', { 
+      port: message.data.port,
+      master_id: incoming_master_id 
+    });
+    this.port = message.data.port;
+    this.settings = message.data.settings;
+    this.master_id = incoming_master_id;
+    
+    // NOTE: Initialize database for read operations in worker process.
+    try {
+      initialize_database();
+      this.log.info('Database initialized in worker process');
+    } catch (error) {
+      this.log.error('Failed to initialize database in worker process', { error: error.message });
+    }
+    
+    this.log.info('Starting server', { port: this.port });
+    this.start_server();
+  }
+  
+  start_server() {
+    this.server = net.createServer((socket) => {
+      this.handle_connection(socket);
+    });
+    
+    this.server.listen(this.port, () => {
+      this.log.info('Server listening', { port: this.port });
+      if (process.connected) {
+        process.send({ type: 'server_ready' });
+      }
+    });
+    
+    this.server.on('error', (error) => {
+      this.log.error('Server error', { error: error.message });
+    });
+  }
+  
+  handle_connection(socket) {
+    const socket_id = `${process.pid}_${Date.now()}_${Math.random()}`;
+    socket.id = socket_id;
+    
+    // NOTE: Create message parser for this socket to handle length-prefixed MessagePack messages.
+    socket.message_parser = create_message_parser();
+    
+    this.connections.set(socket_id, socket);
+    this.connection_count++;
+    this.update_connection_count();
+    
+    socket.on('data', (raw_data) => {
+      this.handle_socket_data(socket, raw_data);
+    });
+    
+    socket.on('end', () => {
+      this.handle_socket_end(socket);
+    });
+    
+    socket.on('error', (error) => {
+      this.log.error('Socket error', { socket_id, error_message: error.message });
+      this.handle_socket_end(socket);
+    });
+  }
+  
+  handle_socket_data(socket, raw_data) {
+    try {
+      // NOTE: Use the message parser to handle length-prefixed MessagePack messages.
+      const messages = socket.message_parser.parse_messages(raw_data);
+      
+      for (const message of messages) {
+        // NOTE: Message is already decoded from MessagePack by the message parser.
+        const parsed_data = message;
+        const op_type = parsed_data?.op || null;
+        
+        if (!op_type) {
+          send_error(socket, { message: 'Missing operation type' });
+          continue;
+        }
+        
+        const is_valid_op_type = this.check_op_type(op_type);
+        
+        if (!is_valid_op_type) {
+          send_error(socket, { message: 'Invalid operation type' });
+          continue;
+        }
+        
+        this.route_operation(socket, op_type, parsed_data?.data || {});
+      }
+      
+    } catch (error) {
+      this.log.error('Data parsing error', { 
+        socket_id: socket.id, 
+        error_message: error.message 
+      });
+      send_error(socket, { message: 'Invalid data format' });
+    }
+  }
+  
+  handle_socket_end(socket) {
+    if (socket.id) {
+      this.connections.delete(socket.id);
+      this.authenticated_clients.delete(socket.id);
+      this.connection_count--;
+      this.update_connection_count();
+    }
+    this.log.info('Client disconnected', { socket_id: socket.id });
+  }
+  
+  
+  check_op_type(op_type = '') {
+    if (!op_type) {
+      return false;
+    }
+    return op_types.includes(op_type);
+  }
+  
+  route_operation(socket, op_type, data) {
+    switch (op_type) {
+      case 'authentication':
+        this.handle_authentication(socket, data);
+        break;
+      case 'setup':
+        this.handle_setup(socket, data);
+        break;
+      case 'find_one':
+      case 'find':
+      case 'get_indexes':
+        // NOTE: Route read operations through master for consistency.
+        this.handle_read_operation(socket, op_type, data);
+        break;
+      case 'create_index':
+      case 'drop_index':
+        // NOTE: Index operations are write operations, route through master.
+        this.handle_write_operation(socket, op_type, data);
+        break;
+      case 'insert_one':
+      case 'update_one':
+      case 'delete_one':
+      case 'bulk_write':
+        this.handle_write_operation(socket, op_type, data);
+        break;
+      case 'ping':
+        handle_ping_operation(socket);
+        break;
+      case 'admin':
+        // NOTE: Use shared operation dispatcher for admin operations.
+        handle_admin_operation(socket, data, this.is_authenticated.bind(this));
+        break;
+      default:
+        send_error(socket, { message: `Unsupported operation: ${op_type}` });
+    }
+  }
+  
+  handle_authentication(socket, data) {
+    if (!this.is_authenticated(socket)) {
+      const auth_id = `${socket.id}_${Date.now()}`;
+      
+      process.send({
+        type: 'auth_request',
+        data: {
+          auth_id,
+          socket_id: socket.id,
+          password: data.password
+        }
+      });
+      
+      this.pending_writes.set(auth_id, { socket, type: 'auth' });
+    } else {
+      send_message(socket, 'Already authenticated');
+    }
+  }
+  
+  handle_setup(socket, data) {
+    // NOTE: Forward setup request to master process.
+    const setup_id = `${socket.id}_${Date.now()}`;
+    
+    process.send({
+      type: 'setup_request',
+      data: {
+        setup_id,
+        socket_id: socket.id
+      }
+    });
+    
+    this.pending_writes.set(setup_id, { socket, type: 'setup' });
+  }
+  
+  
+  handle_read_operation(socket, op_type, data) {
+    if (!this.is_authenticated(socket)) {
+      send_error(socket, { message: 'Authentication required' });
+      return;
+    }
+    
+    const read_id = `${socket.id}_${++this.write_id_counter}`;
+    
+    process.send({
+      type: 'write_request',
+      data: {
+        write_id: read_id,
+        op_type,
+        data,
+        socket_id: socket.id
+      }
+    });
+    
+    this.pending_writes.set(read_id, { socket, type: 'read' });
+  }
+  
+  handle_write_operation(socket, op_type, data) {
+    if (!this.is_authenticated(socket)) {
+      send_error(socket, { message: 'Authentication required' });
+      return;
+    }
+    
+    const write_id = `${socket.id}_${++this.write_id_counter}`;
+    
+    process.send({
+      type: 'write_request',
+      data: {
+        write_id,
+        op_type,
+        data,
+        socket_id: socket.id
+      }
+    });
+    
+    this.pending_writes.set(write_id, { socket, type: 'write' });
+  }
+  
+  
+  
+  
+  handle_write_response(message) {
+    const { write_id, success, result, error } = message.data;
+    const pending_write = this.pending_writes.get(write_id);
+    
+    if (!pending_write) {
+      this.log.warn('No pending write found', { write_id });
+      return;
+    }
+    
+    const { socket } = pending_write;
+    this.pending_writes.delete(write_id);
+    
+    // NOTE: Check if socket is still connected before sending response.
+    if (socket.destroyed || !socket.writable) {
+      this.log.warn('Socket disconnected before response could be sent', { write_id });
+      return;
+    }
+    
+    try {
+      if (success) {
+        // NOTE: Handle different result types properly
+        let response;
+        if (Array.isArray(result)) {
+          // For find operations that return arrays directly
+          response = { ok: 1, documents: result };
+        } else if (result && typeof result === 'object') {
+          // For operations that return objects, spread them
+          response = { ok: 1, ...result };
+        } else {
+          // For simple values
+          response = { ok: 1, result };
+        }
+        const encoded_response = encode_message(response);
+        socket.write(encoded_response);
+      } else {
+        const error_response = { ok: 0, error: error };
+        const encoded_error_response = encode_message(error_response);
+        socket.write(encoded_error_response);
+      }
+    } catch (write_error) {
+      this.log.error('Error sending response to client', { 
+        write_id, 
+        error: write_error.message 
+      });
+    }
+  }
+  
+  handle_auth_response(message) {
+    const { auth_id, success, message: auth_message } = message.data;
+    const pending_auth = this.pending_writes.get(auth_id);
+    
+    if (!pending_auth) {
+      this.log.warn('No pending auth found', { auth_id });
+      return;
+    }
+    
+    const { socket } = pending_auth;
+    this.pending_writes.delete(auth_id);
+    
+    // NOTE: Check if socket is still connected before sending response.
+    if (socket.destroyed || !socket.writable) {
+      this.log.warn('Socket disconnected before auth response could be sent', { auth_id });
+      return;
+    }
+    
+    try {
+      if (success) {
+        this.authenticated_clients.add(socket.id);
+        const auth_response = { 
+          ok: 1, 
+          version: "1.0.0",
+          message: auth_message 
+        };
+        const encoded_auth_response = encode_message(auth_response);
+        socket.write(encoded_auth_response);
+      } else {
+        send_error(socket, { message: auth_message });
+        socket.end();
+      }
+    } catch (write_error) {
+      this.log.error('Error sending auth response to client', { 
+        auth_id, 
+        error: write_error.message 
+      });
+    }
+  }
+  
+  handle_setup_response(message) {
+    const { setup_id, success, password, message: setup_message, error } = message.data;
+    const pending_setup = this.pending_writes.get(setup_id);
+    
+    if (!pending_setup) {
+      this.log.warn('No pending setup found', { setup_id });
+      return;
+    }
+    
+    const { socket } = pending_setup;
+    this.pending_writes.delete(setup_id);
+    
+    if (success) {
+      send_success(socket, {
+        password,
+        message: setup_message
+      });
+    } else {
+      send_error(socket, { message: error });
+    }
+  }
+  
+  handle_write_notification(message) {
+    this.log.info('Received write notification', { 
+      op_type: message.data.op_type,
+      timestamp: message.data.timestamp
+    });
+  }
+  
+  is_authenticated(socket) {
+    return this.authenticated_clients.has(socket.id);
+  }
+  
+  update_connection_count() {
+    if (process.connected) {
+      process.send({
+        type: 'connection_count',
+        data: { count: this.connection_count }
+      });
+    }
+  }
+  
+  send_heartbeat() {
+    if (process.connected) {
+      try {
+        process.send({
+          type: 'heartbeat',
+          data: { timestamp: Date.now() }
+        });
+      } catch (error) {
+        clearInterval(this.heartbeat_interval);
+      }
+    }
+  }
+  
+  async shutdown() {
+    const shutdown_start = Date.now();
+    this.log.info('Initiating graceful shutdown');
+    
+    try {
+      await shutdown_write_queue();
+      this.log.info('Write queue shutdown complete');
+    } catch (error) {
+      this.log.error('Error shutting down write queue', { error: error.message });
+    }
+    
+    // NOTE: Cleanup database before shutting down.
+    try {
+      await cleanup_database();
+      this.log.info('Database cleanup complete');
+    } catch (error) {
+      this.log.error('Error cleaning up database', { error: error.message });
+    }
+    
+    if (this.server) {
+      this.server.close(() => {
+        this.log.info('Server closed');
+      });
+    }
+    
+    for (const [_socket_id, socket] of this.connections) {
+      socket.end();
+    }
+    
+    // NOTE: Use shorter timeout for tests to prevent test timeouts.
+    const shutdown_timeout = process.env.NODE_ENV === 'test' ? 100 : 5000;
+    
+    setTimeout(() => {
+      const shutdown_duration = Date.now() - shutdown_start;
+      this.log.info('Worker shutdown complete', { shutdown_duration_ms: shutdown_duration });
+      process.exit(0);
+    }, shutdown_timeout);
+  }
+}
+
+const worker = new ClusterWorker();