Skip to content

Commit

Permalink
fix(server): delay ws server listen when restart (#12734)
Browse files Browse the repository at this point in the history
  • Loading branch information
sun0day committed Apr 5, 2023
1 parent 1ea38e2 commit abe9274
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 34 deletions.
14 changes: 13 additions & 1 deletion packages/vite/src/node/server/index.ts
Expand Up @@ -338,6 +338,13 @@ export interface ResolvedServerUrls {

export async function createServer(
inlineConfig: InlineConfig = {},
): Promise<ViteDevServer> {
return _createServer(inlineConfig, { ws: true })
}

export async function _createServer(
inlineConfig: InlineConfig = {},
options: { ws: boolean },
): Promise<ViteDevServer> {
const config = await resolveConfig(inlineConfig, 'serve')

Expand Down Expand Up @@ -699,6 +706,9 @@ export async function createServer(
return listen(port, ...args)
}) as any
} else {
if (options.ws) {
ws.listen()
}
await initServer()
}

Expand Down Expand Up @@ -832,7 +842,8 @@ async function restartServer(server: ViteDevServer) {

let newServer = null
try {
newServer = await createServer(inlineConfig)
// delay ws server listen
newServer = await _createServer(inlineConfig, { ws: false })
} catch (err: any) {
server.config.logger.error(err.message, {
timestamp: true,
Expand Down Expand Up @@ -864,6 +875,7 @@ async function restartServer(server: ViteDevServer) {
server.printUrls()
}
} else {
server.ws.listen()
logger.info('server restarted.', { timestamp: true })
}

Expand Down
64 changes: 31 additions & 33 deletions packages/vite/src/node/server/ws.ts
@@ -1,10 +1,10 @@
import type { Server } from 'node:http'
import { STATUS_CODES } from 'node:http'
import { STATUS_CODES, createServer as createHttpServer } from 'node:http'
import type { ServerOptions as HttpsServerOptions } from 'node:https'
import { createServer as createHttpsServer } from 'node:https'
import type { Socket } from 'node:net'
import colors from 'picocolors'
import type { ServerOptions, WebSocket as WebSocketRaw } from 'ws'
import type { WebSocket as WebSocketRaw } from 'ws'
import { WebSocketServer as WebSocketServerRaw } from 'ws'
import type { WebSocket as WebSocketTypes } from 'dep-types/ws'
import type { CustomPayload, ErrorPayload, HMRPayload } from 'types/hmrPayload'
Expand All @@ -20,6 +20,10 @@ export type WebSocketCustomListener<T> = (
) => void

export interface WebSocketServer {
/**
* Listen on port and host
*/
listen(): void
/**
* Get all connected clients.
*/
Expand Down Expand Up @@ -83,7 +87,7 @@ export function createWebSocketServer(
httpsOptions?: HttpsServerOptions,
): WebSocketServer {
let wss: WebSocketServerRaw
let httpsServer: Server | undefined = undefined
let wsHttpServer: Server | undefined = undefined

const hmr = isObject(config.server.hmr) && config.server.hmr
const hmrServer = hmr && hmr.server
Expand All @@ -93,6 +97,8 @@ export function createWebSocketServer(
const wsServer = hmrServer || (portsAreCompatible && server)
const customListeners = new Map<string, Set<WebSocketCustomListener<any>>>()
const clientsMap = new WeakMap<WebSocketRaw, WebSocketClient>()
const port = hmrPort || 24678
const host = (hmr && hmr.host) || undefined

if (wsServer) {
wss = new WebSocketServerRaw({ noServer: true })
Expand All @@ -104,39 +110,28 @@ export function createWebSocketServer(
}
})
} else {
const websocketServerOptions: ServerOptions = {}
const port = hmrPort || 24678
const host = (hmr && hmr.host) || undefined
if (httpsOptions) {
// if we're serving the middlewares over https, the ws library doesn't support automatically creating an https server, so we need to do it ourselves
// create an inline https server and mount the websocket server to it
httpsServer = createHttpsServer(httpsOptions, (req, res) => {
const statusCode = 426
const body = STATUS_CODES[statusCode]
if (!body)
throw new Error(
`No body text found for the ${statusCode} status code`,
)
// http server request handler keeps the same with
// https://github.com/websockets/ws/blob/45e17acea791d865df6b255a55182e9c42e5877a/lib/websocket-server.js#L88-L96
const route = ((_, res) => {
const statusCode = 426
const body = STATUS_CODES[statusCode]
if (!body)
throw new Error(`No body text found for the ${statusCode} status code`)

res.writeHead(statusCode, {
'Content-Length': body.length,
'Content-Type': 'text/plain',
})
res.end(body)
res.writeHead(statusCode, {
'Content-Length': body.length,
'Content-Type': 'text/plain',
})

httpsServer.listen(port, host)
websocketServerOptions.server = httpsServer
res.end(body)
}) as Parameters<typeof createHttpServer>[1]
if (httpsOptions) {
wsHttpServer = createHttpsServer(httpsOptions, route)
} else {
// we don't need to serve over https, just let ws handle its own server
websocketServerOptions.port = port
if (host) {
websocketServerOptions.host = host
}
wsHttpServer = createHttpServer(route)
}

// vite dev server in middleware mode
wss = new WebSocketServerRaw(websocketServerOptions)
// need to call ws listen manually
wss = new WebSocketServerRaw({ server: wsHttpServer })
}

wss.on('connection', (socket) => {
Expand Down Expand Up @@ -210,6 +205,9 @@ export function createWebSocketServer(
let bufferedError: ErrorPayload | null = null

return {
listen: () => {
wsHttpServer?.listen(port, host)
},
on: ((event: string, fn: () => void) => {
if (wsServerEvents.includes(event)) wss.on(event, fn)
else {
Expand Down Expand Up @@ -266,8 +264,8 @@ export function createWebSocketServer(
if (err) {
reject(err)
} else {
if (httpsServer) {
httpsServer.close((err) => {
if (wsHttpServer) {
wsHttpServer.close((err) => {
if (err) {
reject(err)
} else {
Expand Down

0 comments on commit abe9274

Please sign in to comment.