AnimationManager/server/socket-connecting.js

// WebSocket Server
// For matting-server connection

const WebSocket = require('ws');

const WS_PORT = 9527;

// WebSocket server instance
let wss = null;
// Matting server client connection
let mattingClient = null;
// Pending callbacks map
const pendingCallbacks = new Map();

/**
 * Start WebSocket Server
 */
function startWebSocketServer() {
  if (wss) {
    console.log('[WebSocket] Server already running');
    return;
  }

  wss = new WebSocket.Server({ port: WS_PORT });

  wss.on('listening', () => {
    console.log(`[WebSocket] Server listening on port: ${WS_PORT}`);
  });

  wss.on('connection', (ws, req) => {
    const clientIP = req.socket.remoteAddress;
    console.log(`[WebSocket] New connection: ${clientIP}`);
    mattingClient = ws;

    ws.on('message', (data) => {
      try {
        // Binary format: 4-byte filename length + filename + zip data
        if (Buffer.isBuffer(data)) {
          const filenameLen = data.readUInt32BE(0);
          const filename = data.slice(4, 4 + filenameLen).toString('utf-8');
          const zipData = data.slice(4 + filenameLen);

          console.log(`[WebSocket] Received: ${filename}, size: ${zipData.length} bytes`);

          // Find and execute callback
          const callback = pendingCallbacks.get(filename);
          if (callback) {
            pendingCallbacks.delete(filename);
            callback(null, zipData);
          } else {
            console.warn(`[WebSocket] No callback found for: ${filename}`);
          }
        } else {
          console.log(`[WebSocket] Received text: ${data}`);
        }
      } catch (error) {
        console.error('[WebSocket] Message handling error:', error);
      }
    });

    ws.on('close', () => {
      console.log('[WebSocket] Connection closed');
      if (mattingClient === ws) {
        mattingClient = null;
      }
    });

    ws.on('error', (error) => {
      console.error('[WebSocket] Connection error:', error);
    });
  });

  wss.on('error', (error) => {
    console.error('[WebSocket] Server error:', error);
  });
}

/**
 * Check if matting-server is connected
 * @returns {boolean}
 */
function isMattingServerConnected() {
  return mattingClient !== null && mattingClient.readyState === WebSocket.OPEN;
}

/**
 * Send matting task to matting-server
 * @param {string} taskId - Task ID (used as filename)
 * @param {Buffer} zipData - ZIP data buffer
 * @param {Function} callback - Callback function (error, resultZipData)
 * @param {number} timeout - Timeout in milliseconds (default 5 minutes)
 */
function sendMattingTask(taskId, zipData, callback, timeout = 5 * 60 * 1000) {
  if (!isMattingServerConnected()) {
    callback(new Error('matting-server not connected'));
    return;
  }

  try {
    // Binary format: 4-byte filename length + filename + zip data
    const filenameBytes = Buffer.from(taskId, 'utf-8');
    const lenBuffer = Buffer.alloc(4);
    lenBuffer.writeUInt32BE(filenameBytes.length, 0);

    const message = Buffer.concat([lenBuffer, filenameBytes, zipData]);

    // Register callback
    pendingCallbacks.set(taskId, callback);

    // Set timeout
    const timeoutId = setTimeout(() => {
      if (pendingCallbacks.has(taskId)) {
        pendingCallbacks.delete(taskId);
        callback(new Error('VIP matting request timeout'));
      }
    }, timeout);

    // Wrap callback to clear timeout
    const originalCallback = callback;
    pendingCallbacks.set(taskId, (error, data) => {
      clearTimeout(timeoutId);
      originalCallback(error, data);
    });

    // Send message
    mattingClient.send(message);
    console.log(`[WebSocket] Sent task: ${taskId}, size: ${message.length} bytes`);

  } catch (error) {
    pendingCallbacks.delete(taskId);
    callback(error);
  }
}

/**
 * Stop WebSocket Server
 */
function stopWebSocketServer() {
  if (wss) {
    wss.close(() => {
      console.log('[WebSocket] Server stopped');
    });
    wss = null;
    mattingClient = null;
  }
}

// Auto-start WebSocket Server
startWebSocketServer();

module.exports = {
  startWebSocketServer,
  stopWebSocketServer,
  isMattingServerConnected,
  sendMattingTask
};