From 02e91a96ddded6f6d2505f80b45705fcfb77ed74 Mon Sep 17 00:00:00 2001 From: Eugen Rochko Date: Tue, 7 Feb 2017 14:37:12 +0100 Subject: [PATCH] Make streaming API use one pattern-matching redis pubsub connection Refresh timelines when streaming API reconnects in the UI --- .../components/containers/mastodon.jsx | 7 +- app/assets/javascripts/components/stream.jsx | 9 +- streaming/index.js | 94 +++++++++++++------ 3 files changed, 74 insertions(+), 36 deletions(-) diff --git a/app/assets/javascripts/components/containers/mastodon.jsx b/app/assets/javascripts/components/containers/mastodon.jsx index 3b36ce3ef..e23c65121 100644 --- a/app/assets/javascripts/components/containers/mastodon.jsx +++ b/app/assets/javascripts/components/containers/mastodon.jsx @@ -6,7 +6,7 @@ import { deleteFromTimelines, refreshTimeline } from '../actions/timelines'; -import { updateNotifications } from '../actions/notifications'; +import { updateNotifications, refreshNotifications } from '../actions/notifications'; import createBrowserHistory from 'history/lib/createBrowserHistory'; import { applyRouterMiddleware, @@ -80,6 +80,11 @@ const Mastodon = React.createClass({ store.dispatch(updateNotifications(JSON.parse(data.payload), getMessagesForLocale(locale), locale)); break; } + }, + + reconnected () { + store.dispatch(refreshTimeline('home')); + store.dispatch(refreshNotifications()); } }); diff --git a/app/assets/javascripts/components/stream.jsx b/app/assets/javascripts/components/stream.jsx index 0787399f6..392268b23 100644 --- a/app/assets/javascripts/components/stream.jsx +++ b/app/assets/javascripts/components/stream.jsx @@ -10,12 +10,13 @@ const createWebSocketURL = (url) => { return a.href; }; -export default function getStream(accessToken, stream, { connected, received, disconnected }) { +export default function getStream(accessToken, stream, { connected, received, disconnected, reconnected }) { const ws = new WebSocketClient(`${createWebSocketURL(STREAMING_API_BASE_URL)}/api/v1/streaming/?access_token=${accessToken}&stream=${stream}`); - ws.onopen = connected; - ws.onmessage = e => received(JSON.parse(e.data)); - ws.onclose = disconnected; + ws.onopen = connected; + ws.onmessage = e => received(JSON.parse(e.data)); + ws.onclose = disconnected; + ws.onreconnect = reconnected; return ws; }; diff --git a/streaming/index.js b/streaming/index.js index e2e8f943e..125b35bb4 100644 --- a/streaming/index.js +++ b/streaming/index.js @@ -36,6 +36,39 @@ 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 subs = {} + +redisClient.on('pmessage', (_, channel, message) => { + const callbacks = subs[channel] + + log.silly(`New message on channel ${channel}`) + + if (!callbacks) { + return + } + + callbacks.forEach(callback => callback(message)) +}) + +redisClient.psubscribe('timeline:*') + +const subscribe = (channel, callback) => { + log.silly(`Adding listener for ${channel}`) + subs[channel] = subs[channel] || [] + subs[channel].push(callback) +} + +const unsubscribe = (channel, callback) => { + log.silly(`Removing listener for ${channel}`) + subs[channel] = subs[channel].filter(item => item !== callback) +} + const allowCrossDomain = (req, res, next) => { res.header('Access-Control-Allow-Origin', '*') res.header('Access-Control-Allow-Headers', 'Authorization, Accept, Cache-Control') @@ -105,10 +138,10 @@ const errorMiddleware = (err, req, res, next) => { const placeholders = (arr, shift = 0) => arr.map((_, i) => `$${i + 1 + shift}`).join(', '); -const streamFrom = (redisClient, id, req, output, needsFiltering = false) => { +const streamFrom = (id, req, output, attachCloseHandler, needsFiltering = false) => { log.verbose(req.requestId, `Starting stream from ${id} for ${req.accountId}`) - redisClient.on('message', (channel, message) => { + const listener = message => { const { event, payload, queued_at } = JSON.parse(message) const transmit = () => { @@ -149,13 +182,14 @@ const streamFrom = (redisClient, id, req, output, needsFiltering = false) => { } else { transmit() } - }) + } - redisClient.subscribe(id) + subscribe(id, listener) + attachCloseHandler(id, listener) } // Setup stream output to HTTP -const streamToHttp = (req, res, redisClient) => { +const streamToHttp = (req, res) => { res.setHeader('Content-Type', 'text/event-stream') res.setHeader('Transfer-Encoding', 'chunked') @@ -164,7 +198,6 @@ const streamToHttp = (req, res, redisClient) => { req.on('close', () => { log.verbose(req.requestId, `Ending stream for ${req.accountId}`) clearInterval(heartbeat) - redisClient.quit() }) return (event, payload) => { @@ -173,11 +206,17 @@ const streamToHttp = (req, res, redisClient) => { } } +// 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, redisClient) => { +const streamToWs = (req, ws) => { ws.on('close', () => { log.verbose(req.requestId, `Ending stream for ${req.accountId}`) - redisClient.quit() }) return (event, payload) => { @@ -190,12 +229,12 @@ const streamToWs = (req, ws, redisClient) => { } } -// Get new redis connection -const getRedisClient = () => redis.createClient({ - host: process.env.REDIS_HOST || '127.0.0.1', - port: process.env.REDIS_PORT || 6379, - password: process.env.REDIS_PASSWORD -}) +// Setup stream end for WebSockets +const streamWsEnd = ws => (id, listener) => { + ws.on('close', () => { + unsubscribe(id, listener) + }) +} app.use(setRequestId) app.use(allowCrossDomain) @@ -203,28 +242,23 @@ app.use(authenticationMiddleware) app.use(errorMiddleware) app.get('/api/v1/streaming/user', (req, res) => { - const redisClient = getRedisClient() - streamFrom(redisClient, `timeline:${req.accountId}`, req, streamToHttp(req, res, redisClient)) + streamFrom(`timeline:${req.accountId}`, req, streamToHttp(req, res), streamHttpEnd(req)) }) app.get('/api/v1/streaming/public', (req, res) => { - const redisClient = getRedisClient() - streamFrom(redisClient, 'timeline:public', req, streamToHttp(req, res, redisClient), true) + streamFrom('timeline:public', req, streamToHttp(req, res), streamHttpEnd(req), true) }) app.get('/api/v1/streaming/public/local', (req, res) => { - const redisClient = getRedisClient() - streamFrom(redisClient, 'timeline:public:local', req, streamToHttp(req, res, redisClient), true) + streamFrom('timeline:public:local', req, streamToHttp(req, res), streamHttpEnd(req), true) }) app.get('/api/v1/streaming/hashtag', (req, res) => { - const redisClient = getRedisClient() - streamFrom(redisClient, `timeline:hashtag:${req.params.tag}`, req, streamToHttp(req, res, redisClient), true) + streamFrom(`timeline:hashtag:${req.params.tag}`, req, streamToHttp(req, res), streamHttpEnd(req), true) }) app.get('/api/v1/streaming/hashtag/local', (req, res) => { - const redisClient = getRedisClient() - streamFrom(redisClient, `timeline:hashtag:${req.params.tag}:local`, req, streamToHttp(req, res, redisClient), true) + streamFrom(`timeline:hashtag:${req.params.tag}:local`, req, streamToHttp(req, res), streamHttpEnd(req), true) }) wss.on('connection', ws => { @@ -239,23 +273,21 @@ wss.on('connection', ws => { return } - const redisClient = getRedisClient() - switch(location.query.stream) { case 'user': - streamFrom(redisClient, `timeline:${req.accountId}`, req, streamToWs(req, ws, redisClient)) + streamFrom(`timeline:${req.accountId}`, req, streamToWs(req, ws), streamWsEnd(ws)) break; case 'public': - streamFrom(redisClient, 'timeline:public', req, streamToWs(req, ws, redisClient), true) + streamFrom('timeline:public', req, streamToWs(req, ws), streamWsEnd(ws), true) break; case 'public:local': - streamFrom(redisClient, 'timeline:public:local', req, streamToWs(req, ws, redisClient), true) + streamFrom('timeline:public:local', req, streamToWs(req, ws), streamWsEnd(ws), true) break; case 'hashtag': - streamFrom(redisClient, `timeline:hashtag:${location.query.tag}`, req, streamToWs(req, ws, redisClient), true) + streamFrom(`timeline:hashtag:${location.query.tag}`, req, streamToWs(req, ws), streamWsEnd(ws), true) break; case 'hashtag:local': - streamFrom(redisClient, `timeline:hashtag:${location.query.tag}:local`, req, streamToWs(req, ws, redisClient), true) + streamFrom(`timeline:hashtag:${location.query.tag}:local`, req, streamToWs(req, ws), streamWsEnd(ws), true) break; default: ws.close()