SSEEventSourceServlet.java

package com.airbnb.airpal.resources.sse;

import com.airbnb.airpal.api.Job;
import com.airbnb.airpal.api.event.JobEvent;
import com.airbnb.airpal.api.event.JobFinishedEvent;
import com.airbnb.airpal.api.event.JobUpdateEvent;
import com.airbnb.airpal.core.AirpalUser;
import com.airbnb.airpal.core.AirpalUserFactory;
import com.airbnb.airpal.core.AuthorizationUtil;
import com.codahale.metrics.MetricRegistry;
import com.codahale.metrics.Timer;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.collect.Iterables;
import com.google.common.eventbus.EventBus;
import com.google.common.eventbus.Subscribe;
import com.google.common.util.concurrent.RateLimiter;
import com.google.inject.Inject;
import com.google.inject.name.Named;
import lombok.Value;
import lombok.extern.slf4j.Slf4j;
import org.eclipse.jetty.servlets.EventSource;
import org.eclipse.jetty.servlets.EventSourceServlet;

import javax.servlet.http.HttpServletRequest;

import java.util.Collections;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;

import static com.codahale.metrics.MetricRegistry.name;
import static com.google.common.base.Preconditions.checkNotNull;

@Slf4j
public class SSEEventSourceServlet extends EventSourceServlet
{
    private final JobUpdateToSSERelay jobUpdateToSSERelay;
    private final AirpalUserFactory userFactory;

    @Inject
    public SSEEventSourceServlet(ObjectMapper objectMapper,
            EventBus eventBus,
            @Named("sse") ExecutorService executorService,
            MetricRegistry registry,
            AirpalUserFactory userFactory)
    {
        this.jobUpdateToSSERelay = new JobUpdateToSSERelay(objectMapper, executorService, registry);
        this.userFactory = userFactory;
        eventBus.register(jobUpdateToSSERelay);
    }

    @Override
    protected EventSource newEventSource(HttpServletRequest request)
    {
        SSEEventSource eventSource = new SSEEventSource(jobUpdateToSSERelay);
        jobUpdateToSSERelay.addListener(eventSource, userFactory.provide());
        return eventSource;
    }

    static class JobUpdateToSSERelay
    {
        private final ObjectMapper objectMapper;
        private final RateLimiter updateLimiter = RateLimiter.create(15.0);
        private final Set<SSEEventSource> subscribers = Collections.newSetFromMap(new ConcurrentHashMap<SSEEventSource, Boolean>());
        private final Map<SSEEventSource, AirpalUser> eventSourceSubjectMap = new ConcurrentHashMap<>();
        private final ExecutorService executorService;
        private final Timer timer;

        public JobUpdateToSSERelay(ObjectMapper objectMapper, ExecutorService executorService, MetricRegistry registry)
        {
            this.objectMapper = checkNotNull(objectMapper, "objectMapper was null");
            this.executorService = checkNotNull(executorService, "executorService was null");
            this.timer = registry.timer(name(AuthorizedEventBroadcast.class, "authorization"));
        }

        public void addListener(SSEEventSource sseEventSource, AirpalUser subject)
        {
            AirpalUser eventSubject = checkNotNull(subject, "subject was null");
            SSEEventSource eventSource = checkNotNull(sseEventSource, "sseEventSource was null");

            subscribers.add(eventSource);
            eventSourceSubjectMap.put(eventSource, eventSubject);
        }

        public void removeListener(SSEEventSource sseEventSource)
        {
            SSEEventSource eventSource = checkNotNull(sseEventSource, "sseEventSource was null");

            subscribers.remove(eventSource);
            eventSourceSubjectMap.remove(eventSource);
        }

        private void broadcast(JobEvent message)
        {
            try {
                String jsonMessage = objectMapper.writeValueAsString(message);

                for (SSEEventSource subscriber : subscribers) {
                    executorService.submit(
                            new AuthorizedEventBroadcast(subscriber,
                                    eventSourceSubjectMap.get(subscriber),
                                    jsonMessage,
                                    message.getJob(),
                                    timer));
                }
            }
            catch (JsonProcessingException e) {
                log.error("Could not serialize JobEvent as JSON", e);
            }
        }

        @Subscribe
        public void receiveJobUpdate(JobUpdateEvent event) {
            if (updateLimiter.tryAcquire()) {
                broadcast(event);
            }
        }

        @Subscribe
        public void receiveJobFinished(JobFinishedEvent event) {
            broadcast(event);
        }
    }

    @Value
    private static class AuthorizedEventBroadcast implements Runnable
    {
        private final SSEEventSource eventSource;
        private final AirpalUser subject;
        private final String message;
        private final Job job;
        private final Timer timer;

        @Override
        public void run()
        {
            Timer.Context context = timer.time();
            if (Iterables.all(job.getTablesUsed(), new AuthorizationUtil.AuthorizedTablesPredicate(subject))) {
                eventSource.emit(message);
            }
            context.stop();
        }
    }
}