WebSocket 实时推送进度

package.json

@@ -13,6 +13,7 @@
     "cors": "^2.8.5",
     "dotenv": "^16.4.5",
     "express": "^4.19.2",
-    "fluent-ffmpeg": "^2.1.3"
+    "fluent-ffmpeg": "^2.1.3",
+    "ws": "^8.13.0"
   }
 }

public/css/style.css

@@ -447,6 +447,13 @@ header p {
     transform: translateY(-1px);
 }
 
+.play-btn:disabled {
+    opacity: 0.65;
+    background: rgba(37, 99, 235, 0.5);
+    cursor: not-allowed;
+    transform: none;
+}
+
 .play-btn.hidden {
     display: none !important;
 }

public/js/main.js

@@ -15,6 +15,62 @@ document.addEventListener('DOMContentLoaded', () => {
     const progressFill = document.getElementById('progress-fill');
 
     let currentPollInterval = null;
+    let ws = null;
+    let wsConnected = false;
+    let subscribedKey = null;
+    let currentVideoKey = null;
+
+    const sendWsMessage = (message) => {
+        if (ws && ws.readyState === WebSocket.OPEN) {
+            ws.send(JSON.stringify(message));
+        }
+    };
+
+    const subscribeToKey = (key) => {
+        subscribedKey = key;
+        if (wsConnected) {
+            sendWsMessage({ type: 'subscribe', key });
+        }
+    };
+
+    const handleWsMessage = (event) => {
+        try {
+            const message = JSON.parse(event.data);
+            if (message.type === 'progress' && message.key === currentVideoKey) {
+                setProgress(message.progress);
+            }
+            if (message.type === 'ready' && message.key === currentVideoKey) {
+                setProgress({ status: 'Ready to play', percent: 100, details: 'Ready to play' });
+                playHlsStream(message.hlsUrl);
+            }
+        } catch (error) {
+            console.error('WebSocket message parse error:', error);
+        }
+    };
+
+    const connectWebSocket = () => {
+        const protocol = window.location.protocol === 'https:' ? 'wss' : 'ws';
+        ws = new WebSocket(`${protocol}://${window.location.host}`);
+
+        ws.addEventListener('open', () => {
+            wsConnected = true;
+            if (subscribedKey) {
+                sendWsMessage({ type: 'subscribe', key: subscribedKey });
+            }
+        });
+
+        ws.addEventListener('message', handleWsMessage);
+
+        ws.addEventListener('close', () => {
+            wsConnected = false;
+            setTimeout(connectWebSocket, 2000);
+        });
+
+        ws.addEventListener('error', (error) => {
+            console.error('WebSocket error:', error);
+            wsConnected = false;
+        });
+    };
 
     const setProgress = (progress) => {
         if (!progressInfo || !progressText || !progressFill) return;
@@ -161,9 +217,16 @@ document.addEventListener('DOMContentLoaded', () => {
         playerOverlay.classList.add('hidden');
         videoPlayer.classList.add('hidden');
         videoPlayer.pause();
-        if (playBtn) playBtn.classList.add('hidden');
+        if (playBtn) {
+            playBtn.disabled = true;
+            playBtn.textContent = '等待转码完成';
+            playBtn.classList.remove('hidden');
+        }
         resetProgress();
 
+        currentVideoKey = key;
+        subscribeToKey(key);
+
         nowPlaying.classList.remove('hidden');
         currentVideoTitle.textContent = key.split('/').pop();
 
@@ -182,8 +245,9 @@ document.addEventListener('DOMContentLoaded', () => {
 
             if (data.error) throw new Error(data.error);
 
-            // Wait/Poll for HLS playlist to be ready
+            if (!wsConnected) {
-            pollForHlsReady(key, data.hlsUrl);
+                pollForHlsReady(key, data.hlsUrl);
+            }
         } catch (err) {
             console.error(err);
             transcodingOverlay.innerHTML = `<p style="color: #ef4444;">Transcode Failed: ${err.message}</p>`;
@@ -193,7 +257,8 @@ document.addEventListener('DOMContentLoaded', () => {
     // Poll the backend to check if the generated m3u8 file is accessible
     const pollForHlsReady = (key, hlsUrl) => {
         let attempts = 0;
-        const maxAttempts = 60; // 60 seconds max wait for first segment
+        const maxAttempts = 120; // 60 seconds max wait for first segment
+        const pollIntervalMs = 500;
 
         currentPollInterval = setInterval(async () => {
             attempts++;
@@ -216,7 +281,7 @@ document.addEventListener('DOMContentLoaded', () => {
             } catch (err) {
                 console.error("Poll Error:", err);
             }
-        }, 1000);
+        }, pollIntervalMs);
     };
 
     const stopPolling = () => {
@@ -238,13 +303,21 @@ document.addEventListener('DOMContentLoaded', () => {
             hls.loadSource(url);
             hls.attachMedia(videoPlayer);
             hls.on(Hls.Events.MANIFEST_PARSED, () => {
-                playBtn.classList.remove('hidden');
+                if (playBtn) {
+                    playBtn.disabled = false;
+                    playBtn.textContent = 'Play';
+                    playBtn.classList.remove('hidden');
+                }
             });
         } else if (videoPlayer.canPlayType('application/vnd.apple.mpegurl')) {
             // Safari uses native HLS
             videoPlayer.src = url;
             videoPlayer.addEventListener('loadedmetadata', () => {
-                playBtn.classList.remove('hidden');
+                if (playBtn) {
+                    playBtn.disabled = false;
+                    playBtn.textContent = 'Play';
+                    playBtn.classList.remove('hidden');
+                }
             }, { once: true });
         }
     };
@@ -252,6 +325,7 @@ document.addEventListener('DOMContentLoaded', () => {
     // Bind events
     refreshBtn.addEventListener('click', fetchVideos);
 
-    // Initial load
+    // Connect WebSocket and initial load
+    connectWebSocket();
     fetchVideos();
 });

server.js

@@ -3,6 +3,8 @@ const cors = require('cors');
 const dotenv = require('dotenv');
 const fs = require('fs');
 const path = require('path');
+const http = require('http');
+const WebSocket = require('ws');
 const ffmpeg = require('fluent-ffmpeg');
 const { S3Client, ListObjectsV2Command, GetObjectCommand } = require('@aws-sdk/client-s3');
 
@@ -11,6 +13,7 @@ dotenv.config();
 const app = express();
 const PORT = process.env.PORT || 3000;
 const HOST = process.env.HOST || process.env.LISTEN_ADDRESS || '0.0.0.0';
+const server = http.createServer(app);
 
 app.use(cors());
 app.use(express.json());
@@ -29,6 +32,63 @@ const s3Client = new S3Client({
 
 const BUCKET_NAME = process.env.S3_BUCKET_NAME;
 const progressMap = {};
+const wsSubscriptions = new Map();
+
+const addWsClient = (key, ws) => {
+    if (!wsSubscriptions.has(key)) {
+        wsSubscriptions.set(key, new Set());
+    }
+    wsSubscriptions.get(key).add(ws);
+    ws.currentKey = key;
+};
+
+const removeWsClient = (ws) => {
+    const key = ws.currentKey;
+    if (!key) return;
+    const set = wsSubscriptions.get(key);
+    if (!set) return;
+    set.delete(ws);
+    if (set.size === 0) wsSubscriptions.delete(key);
+};
+
+const broadcastWs = (key, payload) => {
+    const clients = wsSubscriptions.get(key);
+    if (!clients) return;
+    const message = JSON.stringify(payload);
+    for (const client of clients) {
+        if (client.readyState === WebSocket.OPEN) {
+            client.send(message);
+        }
+    }
+};
+
+const wss = new WebSocket.Server({ server });
+
+wss.on('connection', (ws) => {
+    ws.on('message', (raw) => {
+        try {
+            const message = JSON.parse(raw.toString());
+            if (message.type === 'subscribe' && typeof message.key === 'string') {
+                if (ws.currentKey && ws.currentKey !== message.key) {
+                    removeWsClient(ws);
+                }
+                addWsClient(message.key, ws);
+
+                const currentProgress = progressMap[message.key];
+                if (currentProgress) {
+                    ws.send(JSON.stringify({ type: 'progress', key: message.key, progress: currentProgress }));
+                    if (currentProgress.status === 'finished' && currentProgress.hlsUrl) {
+                        ws.send(JSON.stringify({ type: 'ready', key: message.key, hlsUrl: currentProgress.hlsUrl }));
+                    }
+                }
+            }
+        } catch (error) {
+            console.error('WebSocket parse error:', error);
+        }
+    });
+
+    ws.on('close', () => removeWsClient(ws));
+});
 
 // Endpoint to list videos in the bucket
 app.get('/api/videos', async (req, res) => {
@@ -88,7 +148,7 @@ app.post('/api/transcode', async (req, res) => {
         const m3u8Path = path.join(outputDir, 'index.m3u8');
         const hlsUrl = `/hls/${progressKey}/index.m3u8`;
 
-        progressMap[progressKey] = { status: 'pending', percent: 0, details: 'Waiting for ffmpeg to start' };
+        progressMap[progressKey] = { status: 'pending', percent: 0, details: 'Waiting for ffmpeg to start', hlsUrl };
 
         // If it already exists, just return the URL
         if (fs.existsSync(m3u8Path)) {
@@ -124,23 +184,31 @@ app.post('/api/transcode', async (req, res) => {
             ])
             .output(m3u8Path)
             .on('progress', (progress) => {
-                progressMap[progressKey] = {
+                const progressState = {
                     status: 'transcoding',
                     percent: Math.min(Math.max(Math.round(progress.percent || 0), 0), 100),
                     frame: progress.frames || null,
                     fps: progress.currentFps || null,
                     bitrate: progress.currentKbps || null,
                     timemark: progress.timemark || null,
-                    details: `Transcoding... ${Math.min(Math.max(Math.round(progress.percent || 0), 0), 100)}%`
+                    details: `Transcoding... ${Math.min(Math.max(Math.round(progress.percent || 0), 0), 100)}%`,
+                    hlsUrl
                 };
+                progressMap[progressKey] = progressState;
+                broadcastWs(progressKey, { type: 'progress', key: progressKey, progress: progressState });
             })
             .on('end', () => {
                 console.log(`Finished transcoding ${key} to HLS`);
-                progressMap[progressKey] = { status: 'finished', percent: 100, details: 'Transcoding complete' };
+                const progressState = { status: 'finished', percent: 100, details: 'Transcoding complete', hlsUrl };
+                progressMap[progressKey] = progressState;
+                broadcastWs(progressKey, { type: 'progress', key: progressKey, progress: progressState });
+                broadcastWs(progressKey, { type: 'ready', key: progressKey, hlsUrl });
             })
             .on('error', (err) => {
                 console.error(`Error transcoding ${key}:`, err);
-                progressMap[progressKey] = { status: 'failed', percent: progressMap[progressKey]?.percent || 0, details: err.message || 'Transcoding failed' };
+                const failedState = { status: 'failed', percent: progressMap[progressKey]?.percent || 0, details: err.message || 'Transcoding failed', hlsUrl };
+                progressMap[progressKey] = failedState;
+                broadcastWs(progressKey, { type: 'progress', key: progressKey, progress: failedState });
             })
             .run();
 
@@ -171,6 +239,6 @@ app.get('/api/status', (req, res) => {
     }
 });
 
-app.listen(PORT, HOST, () => {
+server.listen(PORT, HOST, () => {
     console.log(`Server running on http://${HOST}:${PORT}`);
 });