Skip to content
Snippets Groups Projects
Commit 64e1d510 authored by geta6's avatar geta6 Committed by Eugen
Browse files

Improve streaming server with cluster (#1970)

parent 7a508672
No related branches found
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
.env.production.sample
+ 4
− 0
View file @ 64e1d510
......@@ -81,3 +81,7 @@ SMTP_FROM_ADDRESS=notifications@example.com
# Advanced settings
# If you need to use pgBouncer, you need to disable prepared statements:
# PREPARED_STATEMENTS=false
# Cluster number setting for streaming API server.
# If you comment out following line, cluster number will be `numOfCpuCores - 1`.
STREAMING_CLUSTER_NUM=1
streaming/index.js
+ 244
− 224
View file @ 64e1d510
import os from 'os';
import cluster from 'cluster';
import dotenv from 'dotenv'
import express from 'express'
import http from 'http'
......@@ -14,300 +16,318 @@ dotenv.config({
path: env === 'production' ? '.env.production' : '.env'
})
const pgConfigs = {
development: {
database: 'mastodon_development',
host: '/var/run/postgresql',
max: 10
},
production: {
user: process.env.DB_USER || 'mastodon',
password: process.env.DB_PASS || '',
database: process.env.DB_NAME || 'mastodon_production',
host: process.env.DB_HOST || 'localhost',
port: process.env.DB_PORT || 5432,
max: 10
if (cluster.isMaster) {
// cluster master
const core = +process.env.STREAMING_CLUSTER_NUM || (env === 'development' ? 1 : os.cpus().length - 1)
const fork = () => {
const worker = cluster.fork();
worker.on('exit', (code, signal) => {
log.error(`Worker died with exit code ${code}, signal ${signal} received.`);
setTimeout(() => fork(), 0);
});
};
for (let i = 0; i < core; i++) fork();
log.info(`Starting streaming API server master with ${core} workers`)
} else {
// cluster worker
const pgConfigs = {
development: {
database: 'mastodon_development',
host: '/var/run/postgresql',
max: 10
},
production: {
user: process.env.DB_USER || 'mastodon',
password: process.env.DB_PASS || '',
database: process.env.DB_NAME || 'mastodon_production',
host: process.env.DB_HOST || 'localhost',
port: process.env.DB_PORT || 5432,
max: 10
}
}
}
const app = express()
const pgPool = new pg.Pool(pgConfigs[env])
const server = http.createServer(app)
const wss = new WebSocket.Server({ server })
const redisClient = redis.createClient({
host: process.env.REDIS_HOST || '127.0.0.1',
port: process.env.REDIS_PORT || 6379,
password: process.env.REDIS_PASSWORD
})
const app = express()
const pgPool = new pg.Pool(pgConfigs[env])
const server = http.createServer(app)
const wss = new WebSocket.Server({ server })
const subs = {}
redisClient.on('pmessage', (_, channel, message) => {
const callbacks = subs[channel]
const redisClient = redis.createClient({
host: process.env.REDIS_HOST || '127.0.0.1',
port: process.env.REDIS_PORT || 6379,
password: process.env.REDIS_PASSWORD
})
log.silly(`New message on channel ${channel}`)
const subs = {}
if (!callbacks) {
return
}
redisClient.on('pmessage', (_, channel, message) => {
const callbacks = subs[channel]
callbacks.forEach(callback => callback(message))
})
log.silly(`New message on channel ${channel}`)
redisClient.psubscribe('timeline:*')
if (!callbacks) {
return
}
const subscribe = (channel, callback) => {
log.silly(`Adding listener for ${channel}`)
subs[channel] = subs[channel] || []
subs[channel].push(callback)
}
callbacks.forEach(callback => callback(message))
})
const unsubscribe = (channel, callback) => {
log.silly(`Removing listener for ${channel}`)
subs[channel] = subs[channel].filter(item => item !== callback)
}
redisClient.psubscribe('timeline:*')
const allowCrossDomain = (req, res, next) => {
res.header('Access-Control-Allow-Origin', '*')
res.header('Access-Control-Allow-Headers', 'Authorization, Accept, Cache-Control')
res.header('Access-Control-Allow-Methods', 'GET, OPTIONS')
const subscribe = (channel, callback) => {
log.silly(`Adding listener for ${channel}`)
subs[channel] = subs[channel] || []
subs[channel].push(callback)
}
next()
}
const unsubscribe = (channel, callback) => {
log.silly(`Removing listener for ${channel}`)
subs[channel] = subs[channel].filter(item => item !== callback)
}
const setRequestId = (req, res, next) => {
req.requestId = uuid.v4()
res.header('X-Request-Id', req.requestId)
const allowCrossDomain = (req, res, next) => {
res.header('Access-Control-Allow-Origin', '*')
res.header('Access-Control-Allow-Headers', 'Authorization, Accept, Cache-Control')
res.header('Access-Control-Allow-Methods', 'GET, OPTIONS')
next()
}
next()
}
const accountFromToken = (token, req, next) => {
pgPool.connect((err, client, done) => {
if (err) {
next(err)
return
}
const setRequestId = (req, res, next) => {
req.requestId = uuid.v4()
res.header('X-Request-Id', req.requestId)
client.query('SELECT oauth_access_tokens.resource_owner_id, users.account_id FROM oauth_access_tokens INNER JOIN users ON oauth_access_tokens.resource_owner_id = users.id WHERE oauth_access_tokens.token = $1 LIMIT 1', [token], (err, result) => {
done()
next()
}
const accountFromToken = (token, req, next) => {
pgPool.connect((err, client, done) => {
if (err) {
next(err)
return
}
if (result.rows.length === 0) {
err = new Error('Invalid access token')
err.statusCode = 401
next(err)
return
}
client.query('SELECT oauth_access_tokens.resource_owner_id, users.account_id FROM oauth_access_tokens INNER JOIN users ON oauth_access_tokens.resource_owner_id = users.id WHERE oauth_access_tokens.token = $1 LIMIT 1', [token], (err, result) => {
done()
req.accountId = result.rows[0].account_id
next()
})
})
}
if (err) {
next(err)
return
}
const authenticationMiddleware = (req, res, next) => {
if (req.method === 'OPTIONS') {
next()
return
}
if (result.rows.length === 0) {
err = new Error('Invalid access token')
err.statusCode = 401
const authorization = req.get('Authorization')
next(err)
return
}
if (!authorization) {
const err = new Error('Missing access token')
err.statusCode = 401
req.accountId = result.rows[0].account_id
next(err)
return
next()
})
})
}
const token = authorization.replace(/^Bearer /, '')
const authenticationMiddleware = (req, res, next) => {
if (req.method === 'OPTIONS') {
next()
return
}
accountFromToken(token, req, next)
}
const authorization = req.get('Authorization')
const errorMiddleware = (err, req, res, next) => {
log.error(req.requestId, err)
res.writeHead(err.statusCode || 500, { 'Content-Type': 'application/json' })
res.end(JSON.stringify({ error: err.statusCode ? `${err}` : 'An unexpected error occurred' }))
}
if (!authorization) {
const err = new Error('Missing access token')
err.statusCode = 401
const placeholders = (arr, shift = 0) => arr.map((_, i) => `$${i + 1 + shift}`).join(', ');
next(err)
return
}
const streamFrom = (id, req, output, attachCloseHandler, needsFiltering = false) => {
log.verbose(req.requestId, `Starting stream from ${id} for ${req.accountId}`)
const token = authorization.replace(/^Bearer /, '')
const listener = message => {
const { event, payload, queued_at } = JSON.parse(message)
accountFromToken(token, req, next)
}
const transmit = () => {
const now = new Date().getTime()
const delta = now - queued_at;
const errorMiddleware = (err, req, res, next) => {
log.error(req.requestId, err)
res.writeHead(err.statusCode || 500, { 'Content-Type': 'application/json' })
res.end(JSON.stringify({ error: err.statusCode ? `${err}` : 'An unexpected error occurred' }))
}
log.silly(req.requestId, `Transmitting for ${req.accountId}: ${event} ${payload} Delay: ${delta}ms`)
output(event, payload)
}
const placeholders = (arr, shift = 0) => arr.map((_, i) => `$${i + 1 + shift}`).join(', ');
// Only messages that may require filtering are statuses, since notifications
// are already personalized and deletes do not matter
if (needsFiltering && event === 'update') {
pgPool.connect((err, client, done) => {
if (err) {
log.error(err)
return
}
const streamFrom = (id, req, output, attachCloseHandler, needsFiltering = false) => {
log.verbose(req.requestId, `Starting stream from ${id} for ${req.accountId}`)
const listener = message => {
const { event, payload, queued_at } = JSON.parse(message)
const unpackedPayload = JSON.parse(payload)
const targetAccountIds = [unpackedPayload.account.id].concat(unpackedPayload.mentions.map(item => item.id)).concat(unpackedPayload.reblog ? [unpackedPayload.reblog.account.id] : [])
const transmit = () => {
const now = new Date().getTime()
const delta = now - queued_at;
client.query(`SELECT target_account_id FROM blocks WHERE account_id = $1 AND target_account_id IN (${placeholders(targetAccountIds, 1)}) UNION SELECT target_account_id FROM mutes WHERE account_id = $1 AND target_account_id IN (${placeholders(targetAccountIds, 1)})`, [req.accountId].concat(targetAccountIds), (err, result) => {
done()
log.silly(req.requestId, `Transmitting for ${req.accountId}: ${event} ${payload} Delay: ${delta}ms`)
output(event, payload)
}
// Only messages that may require filtering are statuses, since notifications
// are already personalized and deletes do not matter
if (needsFiltering && event === 'update') {
pgPool.connect((err, client, done) => {
if (err) {
log.error(err)
return
}
if (result.rows.length > 0) {
return
}
const unpackedPayload = JSON.parse(payload)
const targetAccountIds = [unpackedPayload.account.id].concat(unpackedPayload.mentions.map(item => item.id)).concat(unpackedPayload.reblog ? [unpackedPayload.reblog.account.id] : [])
client.query(`SELECT target_account_id FROM blocks WHERE account_id = $1 AND target_account_id IN (${placeholders(targetAccountIds, 1)}) UNION SELECT target_account_id FROM mutes WHERE account_id = $1 AND target_account_id IN (${placeholders(targetAccountIds, 1)})`, [req.accountId].concat(targetAccountIds), (err, result) => {
done()
if (err) {
log.error(err)
return
}
transmit()
if (result.rows.length > 0) {
return
}
transmit()
})
})
})
} else {
transmit()
} else {
transmit()
}
}
subscribe(id, listener)
attachCloseHandler(id, listener)
}
subscribe(id, listener)
attachCloseHandler(id, listener)
}
// Setup stream output to HTTP
const streamToHttp = (req, res) => {
res.setHeader('Content-Type', 'text/event-stream')
res.setHeader('Transfer-Encoding', 'chunked')
// Setup stream output to HTTP
const streamToHttp = (req, res) => {
res.setHeader('Content-Type', 'text/event-stream')
res.setHeader('Transfer-Encoding', 'chunked')
const heartbeat = setInterval(() => res.write(':thump\n'), 15000)
const heartbeat = setInterval(() => res.write(':thump\n'), 15000)
req.on('close', () => {
log.verbose(req.requestId, `Ending stream for ${req.accountId}`)
clearInterval(heartbeat)
})
req.on('close', () => {
log.verbose(req.requestId, `Ending stream for ${req.accountId}`)
clearInterval(heartbeat)
})
return (event, payload) => {
res.write(`event: ${event}\n`)
res.write(`data: ${payload}\n\n`)
}
}
return (event, payload) => {
res.write(`event: ${event}\n`)
res.write(`data: ${payload}\n\n`)
// Setup stream end for HTTP
const streamHttpEnd = req => (id, listener) => {
req.on('close', () => {
unsubscribe(id, listener)
})
}
}
// Setup stream end for HTTP
const streamHttpEnd = req => (id, listener) => {
req.on('close', () => {
unsubscribe(id, listener)
})
}
// Setup stream output to WebSockets
const streamToWs = (req, ws) => {
const heartbeat = setInterval(() => ws.ping(), 15000)
// Setup stream output to WebSockets
const streamToWs = (req, ws) => {
const heartbeat = setInterval(() => ws.ping(), 15000)
ws.on('close', () => {
log.verbose(req.requestId, `Ending stream for ${req.accountId}`)
clearInterval(heartbeat)
})
ws.on('close', () => {
log.verbose(req.requestId, `Ending stream for ${req.accountId}`)
clearInterval(heartbeat)
})
return (event, payload) => {
if (ws.readyState !== ws.OPEN) {
log.error(req.requestId, 'Tried writing to closed socket')
return
}
return (event, payload) => {
if (ws.readyState !== ws.OPEN) {
log.error(req.requestId, 'Tried writing to closed socket')
return
ws.send(JSON.stringify({ event, payload }))
}
ws.send(JSON.stringify({ event, payload }))
}
}
// Setup stream end for WebSockets
const streamWsEnd = ws => (id, listener) => {
ws.on('close', () => {
unsubscribe(id, listener)
})
// Setup stream end for WebSockets
const streamWsEnd = ws => (id, listener) => {
ws.on('close', () => {
unsubscribe(id, listener)
})
ws.on('error', e => {
unsubscribe(id, listener)
})
}
ws.on('error', e => {
unsubscribe(id, listener)
})
}
app.use(setRequestId)
app.use(allowCrossDomain)
app.use(authenticationMiddleware)
app.use(errorMiddleware)
app.use(setRequestId)
app.use(allowCrossDomain)
app.use(authenticationMiddleware)
app.use(errorMiddleware)
app.get('/api/v1/streaming/user', (req, res) => {
streamFrom(`timeline:${req.accountId}`, req, streamToHttp(req, res), streamHttpEnd(req))
})
app.get('/api/v1/streaming/user', (req, res) => {
streamFrom(`timeline:${req.accountId}`, req, streamToHttp(req, res), streamHttpEnd(req))
})
app.get('/api/v1/streaming/public', (req, res) => {
streamFrom('timeline:public', req, streamToHttp(req, res), streamHttpEnd(req), true)
})
app.get('/api/v1/streaming/public', (req, res) => {
streamFrom('timeline:public', req, streamToHttp(req, res), streamHttpEnd(req), true)
})
app.get('/api/v1/streaming/public/local', (req, res) => {
streamFrom('timeline:public:local', req, streamToHttp(req, res), streamHttpEnd(req), true)
})
app.get('/api/v1/streaming/public/local', (req, res) => {
streamFrom('timeline:public:local', req, streamToHttp(req, res), streamHttpEnd(req), true)
})
app.get('/api/v1/streaming/hashtag', (req, res) => {
streamFrom(`timeline:hashtag:${req.params.tag}`, req, streamToHttp(req, res), streamHttpEnd(req), true)
})
app.get('/api/v1/streaming/hashtag', (req, res) => {
streamFrom(`timeline:hashtag:${req.params.tag}`, req, streamToHttp(req, res), streamHttpEnd(req), true)
})
app.get('/api/v1/streaming/hashtag/local', (req, res) => {
streamFrom(`timeline:hashtag:${req.params.tag}:local`, req, streamToHttp(req, res), streamHttpEnd(req), true)
})
app.get('/api/v1/streaming/hashtag/local', (req, res) => {
streamFrom(`timeline:hashtag:${req.params.tag}:local`, req, streamToHttp(req, res), streamHttpEnd(req), true)
})
wss.on('connection', ws => {
const location = url.parse(ws.upgradeReq.url, true)
const token = location.query.access_token
const req = { requestId: uuid.v4() }
wss.on('connection', ws => {
const location = url.parse(ws.upgradeReq.url, true)
const token = location.query.access_token
const req = { requestId: uuid.v4() }
accountFromToken(token, req, err => {
if (err) {
log.error(req.requestId, err)
ws.close()
return
}
accountFromToken(token, req, err => {
if (err) {
log.error(req.requestId, err)
ws.close()
return
}
switch(location.query.stream) {
case 'user':
streamFrom(`timeline:${req.accountId}`, req, streamToWs(req, ws), streamWsEnd(ws))
break;
case 'public':
streamFrom('timeline:public', req, streamToWs(req, ws), streamWsEnd(ws), true)
break;
case 'public:local':
streamFrom('timeline:public:local', req, streamToWs(req, ws), streamWsEnd(ws), true)
break;
case 'hashtag':
streamFrom(`timeline:hashtag:${location.query.tag}`, req, streamToWs(req, ws), streamWsEnd(ws), true)
break;
case 'hashtag:local':
streamFrom(`timeline:hashtag:${location.query.tag}:local`, req, streamToWs(req, ws), streamWsEnd(ws), true)
break;
default:
ws.close()
}
switch(location.query.stream) {
case 'user':
streamFrom(`timeline:${req.accountId}`, req, streamToWs(req, ws), streamWsEnd(ws))
break;
case 'public':
streamFrom('timeline:public', req, streamToWs(req, ws), streamWsEnd(ws), true)
break;
case 'public:local':
streamFrom('timeline:public:local', req, streamToWs(req, ws), streamWsEnd(ws), true)
break;
case 'hashtag':
streamFrom(`timeline:hashtag:${location.query.tag}`, req, streamToWs(req, ws), streamWsEnd(ws), true)
break;
case 'hashtag:local':
streamFrom(`timeline:hashtag:${location.query.tag}:local`, req, streamToWs(req, ws), streamWsEnd(ws), true)
break;
default:
ws.close()
}
})
})
})
server.listen(process.env.PORT || 4000, () => {
log.level = process.env.LOG_LEVEL || 'verbose'
log.info(`Starting streaming API server on port ${server.address().port}`)
})
server.listen(process.env.PORT || 4000, () => {
log.level = process.env.LOG_LEVEL || 'verbose'
log.info(`Starting streaming API server worker on port ${server.address().port}`)
})
}
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment