PushManagerImp.java

package io.split.engine.common;

import com.google.common.annotations.VisibleForTesting;
import io.split.client.interceptors.FlagSetsFilter;
import io.split.engine.experiments.RuleBasedSegmentParser;
import io.split.engine.experiments.SplitParser;
import io.split.engine.sse.AuthApiClient;
import io.split.engine.sse.AuthApiClientImp;
import io.split.engine.sse.EventSourceClient;
import io.split.engine.sse.EventSourceClientImp;
import io.split.engine.sse.PushStatusTracker;
import io.split.engine.sse.PushStatusTrackerImp;
import io.split.engine.sse.client.SSEClient;
import io.split.engine.sse.dtos.AuthenticationResponse;
import io.split.engine.sse.dtos.SegmentQueueDto;
import io.split.engine.sse.workers.SegmentsWorkerImp;
import io.split.engine.sse.workers.FeatureFlagsWorker;
import io.split.engine.sse.workers.FeatureFlagWorkerImp;
import io.split.engine.sse.workers.Worker;

import io.split.storages.RuleBasedSegmentCache;
import io.split.storages.SplitCacheProducer;
import io.split.telemetry.domain.StreamingEvent;
import io.split.telemetry.domain.enums.StreamEventsEnum;
import io.split.telemetry.storage.TelemetryRuntimeProducer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.concurrent.Future;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

import static com.google.common.base.Preconditions.checkNotNull;
import static io.split.client.utils.SplitExecutorFactory.buildSingleThreadScheduledExecutor;

public class PushManagerImp implements PushManager {
    private static final Logger _log = LoggerFactory.getLogger(PushManager.class);

    private final AuthApiClient _authApiClient;
    private final EventSourceClient _eventSourceClient;
    private final FeatureFlagsWorker _featureFlagsWorker;
    private final Worker<SegmentQueueDto> _segmentWorker;
    private final PushStatusTracker _pushStatusTracker;
    private static final Lock lock = new ReentrantLock();

    private Future<?> _nextTokenRefreshTask;
    private final ScheduledExecutorService _scheduledExecutorService;
    private AtomicLong _expirationTime;
    private final TelemetryRuntimeProducer _telemetryRuntimeProducer;
    private final int _streamingTokenRefreshRate;

    @VisibleForTesting
    /* package private */ PushManagerImp(AuthApiClient authApiClient,
                                         EventSourceClient eventSourceClient,
                                         FeatureFlagsWorker featureFlagsWorker,
                                         Worker<SegmentQueueDto> segmentWorker,
                                         PushStatusTracker pushStatusTracker,
                                         TelemetryRuntimeProducer telemetryRuntimeProducer,
                                         ThreadFactory threadFactory,
                                         int streamingTokenRefreshRate) {

        _authApiClient = checkNotNull(authApiClient);
        _eventSourceClient = checkNotNull(eventSourceClient);
        _featureFlagsWorker = featureFlagsWorker;
        _segmentWorker = segmentWorker;
        _pushStatusTracker = pushStatusTracker;
        _expirationTime = new AtomicLong();
        _scheduledExecutorService = buildSingleThreadScheduledExecutor(threadFactory, "Split-SSERefreshToken-%d");
        _telemetryRuntimeProducer = checkNotNull(telemetryRuntimeProducer);
        _streamingTokenRefreshRate = streamingTokenRefreshRate;
    }

    public static PushManagerImp build(Synchronizer synchronizer,
                                       String streamingUrl,
                                       String authUrl,
                                       SplitAPI splitAPI,
                                       LinkedBlockingQueue<PushManager.Status> statusMessages,
                                       TelemetryRuntimeProducer telemetryRuntimeProducer,
                                       ThreadFactory threadFactory,
                                       SplitParser splitParser,
                                       SplitCacheProducer splitCacheProducer,
                                       FlagSetsFilter flagSetsFilter,
                                       RuleBasedSegmentCache ruleBasedSegmentCache,
                                       RuleBasedSegmentParser ruleBasedSegmentParser,
                                       int streamingTokenRefreshRate) {
        FeatureFlagsWorker featureFlagsWorker = new FeatureFlagWorkerImp(synchronizer, splitParser, ruleBasedSegmentParser, splitCacheProducer,
                ruleBasedSegmentCache, telemetryRuntimeProducer, flagSetsFilter);
        Worker<SegmentQueueDto> segmentWorker = new SegmentsWorkerImp(synchronizer);
        PushStatusTracker pushStatusTracker = new PushStatusTrackerImp(statusMessages, telemetryRuntimeProducer);

        return new PushManagerImp(new AuthApiClientImp(authUrl, splitAPI.getHttpClient(), telemetryRuntimeProducer),
                EventSourceClientImp.build(streamingUrl, featureFlagsWorker, segmentWorker, pushStatusTracker, splitAPI.getSseHttpClient(),
                        telemetryRuntimeProducer, threadFactory, splitAPI.getRequestDecorator()),
                featureFlagsWorker,
                segmentWorker,
                pushStatusTracker,
                telemetryRuntimeProducer,
                threadFactory,
                streamingTokenRefreshRate);
    }

    @Override
    public void start() {
        _log.debug("#1 - Start PushManagerImp");
        try {
            lock.lock();
            AuthenticationResponse response = _authApiClient.Authenticate();
            _log.debug(String.format("Auth service response pushEnabled: %s", response.isPushEnabled()));
            if (response.isPushEnabled() && startSse(response.getToken(), response.getChannels())) {
                _log.debug("#2 - PushManagerImp connected");
                _expirationTime.set(_streamingTokenRefreshRate);
                _telemetryRuntimeProducer.recordStreamingEvents(new StreamingEvent(StreamEventsEnum.TOKEN_REFRESH.getType(),
                        response.getExpiration(), System.currentTimeMillis()));
                return;
            }

            _log.debug("#3 - PushManagerImp error");
            cleanUpResources();
            if (response.isRetry()) {
                _log.debug(String.format("Handling retry error response"));
                _pushStatusTracker.handleSseStatus(SSEClient.StatusMessage.RETRYABLE_ERROR);
            } else {
                _log.debug(String.format("Auth service response is disabled: %s", response.getToken()));
                _pushStatusTracker.forcePushDisable();
            }
            _log.debug("#4 - PushManagerImp error");
        } catch (Exception e) {
            _log.debug("Exception in PushManager start: " + e.getMessage());
        } finally {
            lock.unlock();
        }
    }

    @Override
    public void stop() {
        _log.debug("#1 - Stopping PushManagerImp");
        try {
            lock.lock();
            cleanUpResources();
            _log.debug("#2 - Stopped PushManagerImp");
        } finally {
            lock.unlock();
        }
    }

    @Override
    public void scheduleConnectionReset() {
        _log.debug(String.format("scheduleNextTokenRefresh in %s SECONDS", _expirationTime));
        _nextTokenRefreshTask = _scheduledExecutorService.schedule(() -> {
            _log.debug("#1 - Starting scheduleNextTokenRefresh ...");
            stop();
            _log.debug("#2 - Finished to stop all streaming engine");
            start();
            _log.debug("#3 - Finished to start streaming connection");
        }, _expirationTime.get(), TimeUnit.SECONDS);
    }

    private boolean startSse(String token, String channels) {
        try {
            _log.debug("SSE Handler starting ...");
            return _eventSourceClient.start(channels, token);
        } catch (Exception e) {
            _log.debug("Exception in SSE Handler start: " + e.getMessage());
            return false;
        }
    }

    @Override
    public void startWorkers() {
        try {
            _featureFlagsWorker.start();
            _segmentWorker.start();
        } catch (Exception e) {
            _log.debug("Exception in starting workers: " + e.getMessage());
        }
    }

    @Override
    public void stopWorkers() {
        try {
            _featureFlagsWorker.stop();
            _segmentWorker.stop();
        } catch (Exception e) {
            _log.debug("Exception in stopping workers: " + e.getMessage());
        }
    }

    private void cleanUpResources() {
        _log.debug("Starting cleanUpResources - #1")
        _eventSourceClient.stop();
        _log.debug("cleanUpResources - #2")
        stopWorkers();
        if (_nextTokenRefreshTask != null) {
            _log.debug("Cancel nextTokenRefreshTask");
            _nextTokenRefreshTask.cancel(false);
            _log.debug("Finished cleanUpResources - #3 - Finished cancel nextTokenRefreshTask")
        }
        _log.debug("Finished cleanUpResources - #4")
    }
}