com.github.twitch4j.eventsub.socket.TwitchEventSocketPool Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of twitch4j-eventsub-websocket Show documentation
EventSub WebSocket dependency
There is a newer version: 1.23.0
Show newest version
package com.github.twitch4j.eventsub.socket;

import com.github.philippheuer.credentialmanager.domain.OAuth2Credential;
import com.github.philippheuer.events4j.core.EventManager;
import com.github.philippheuer.events4j.simple.SimpleEventHandler;
import com.github.twitch4j.auth.providers.TwitchIdentityProvider;
import com.github.twitch4j.common.pool.SubscriptionConnectionPool;
import com.github.twitch4j.common.util.EventManagerUtils;
import com.github.twitch4j.eventsub.EventSubSubscription;
import com.github.twitch4j.helix.TwitchHelix;
import com.github.twitch4j.helix.TwitchHelixBuilder;
import lombok.Builder;
import lombok.Getter;
import lombok.Synchronized;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.RandomStringUtils;
import org.apache.commons.lang3.StringUtils;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.Collection;
import java.util.Collections;
import java.util.Comparator;
import java.util.LinkedList;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.UnaryOperator;

/**
 * A pool for EventSub websocket subscriptions across multiple users.
 */
@Slf4j
@Builder
public final class TwitchEventSocketPool implements IEventSubSocket {

    private final String threadPrefix = "twitch4j-multi-pool-" + RandomStringUtils.random(4, true, true) + "-eventsub-ws-";

    /**
     * The default {@link EventManager} for this connection pool, if specified.
     */
    @Getter
    @Builder.Default
    private final EventManager eventManager = EventManagerUtils.initializeEventManager(SimpleEventHandler.class);

    /**
     * The {@link ScheduledThreadPoolExecutor} to be used by connections in this pool, if specified.
     */
    @Nullable
    private final ScheduledThreadPoolExecutor executor;

    /**
     * The {@link TwitchIdentityProvider} to enrich credentials.
     */
    @NotNull
    @Builder.Default
    private final TwitchIdentityProvider identityProvider = new TwitchIdentityProvider(null, null, null);

    /**
     * The base url for websocket connections.
     *
     * @see TwitchEventSocket#WEB_SOCKET_SERVER
     */
    @NotNull
    @Builder.Default
    private final String baseUrl = TwitchEventSocket.WEB_SOCKET_SERVER;

    /**
     * The {@link TwitchHelix} instance for creating eventsub subscriptions in the official API.
     */
    @Nullable
    @Builder.Default
    private TwitchHelix helix = TwitchHelixBuilder.builder().build();

    /**
     * The maximum number of eventsub subscriptions that a single user_id can have.
     */
    @Builder.Default
    private int maxSubscriptionsPerUser = TwitchEventSocket.MAX_SUBSCRIPTIONS_PER_SOCKET * 3; // imposed by twitch

    /**
     * Further configuration that should be applied to the builder when creating new EventSocket (single-user) pools.
     */
    @Builder.Default
    private final UnaryOperator> advancedConfiguration = b -> b;

    /**
     * A mapping of user_id's to their individual eventsocket pools.
     */
    private final Map poolByUserId = new ConcurrentHashMap<>();

    /**
     * A mapping of eventsub subscriptions to which individual pool contains it.
     */
    private final Map poolBySub = new ConcurrentHashMap<>();

    @Override
    public void connect() {
        // no-op
    }

    @Override
    public void disconnect() {
        poolByUserId.values().forEach(IEventSubSocket::disconnect);
    }

    @Override
    public void reconnect() {
        poolByUserId.values().forEach(IEventSubSocket::reconnect);
    }

    @Override
    @Synchronized
    public boolean register(OAuth2Credential credential, EventSubSubscription sub) {
        OAuth2Credential token = credential != null ? credential : getDefaultToken();
        if (token == null) return false;

        String userId = getUserId(token);
        if (userId == null) return false;

        SubscriptionWrapper wrapped = SubscriptionWrapper.wrap(sub);

        if (poolBySub.containsKey(wrapped))
            return false;

        TwitchSingleUserEventSocketPool pool = poolByUserId.computeIfAbsent(userId,
            id -> advancedConfiguration.apply(
                TwitchSingleUserEventSocketPool.builder()
                    .baseUrl(baseUrl)
                    .defaultToken(token)
                    .eventManager(eventManager)
                    .helix(helix)
                    .executor(() -> executor)
            ).build()
        );

        if (pool.numSubscriptions() >= maxSubscriptionsPerUser) {
            log.debug("Skipping eventsocket subscription registration because pool is already at capacity for user {}: {}", userId, sub);
            return false;
        }

        return pool.register(token, sub) && poolBySub.put(wrapped, pool) == null;
    }

    @Override
    @Synchronized
    public boolean unregister(EventSubSubscription sub) {
        SubscriptionWrapper wrapped = SubscriptionWrapper.wrap(sub);
        TwitchSingleUserEventSocketPool pool = poolBySub.get(wrapped);
        if (pool == null) return false;

        Boolean unsubscribe = pool.unsubscribe(wrapped);

        // cleanup if we removed the last subscription
        if (pool.numSubscriptions() <= 0) {
            poolByUserId.entrySet().stream()
                .filter(e -> e.getValue() == pool)
                .map(Map.Entry::getKey)
                .findAny()
                .ifPresent(userId -> {
                    AtomicBoolean close = new AtomicBoolean();

                    // noinspection resource
                    poolByUserId.computeIfPresent(userId, (k, v) -> {
                        if (v.numSubscriptions() <= 0) {
                            close.set(true);
                            return null; // remove mapping
                        }
                        return v;
                    });

                    if (close.get())
                        pool.close();
                });
        }

        // noinspection resource
        return unsubscribe != null && unsubscribe && poolBySub.remove(wrapped) != null;
    }

    @Override
    public Collection getSubscriptions() {
        return Collections.unmodifiableSet(poolBySub.keySet());
    }

    @Override
    @Synchronized
    public void close() throws Exception {
        poolBySub.clear();
        Collection pools = new LinkedList<>();
        poolByUserId.values().removeIf(pools::add);
        pools.forEach(SubscriptionConnectionPool::close);
    }

    @Nullable
    @Override
    public OAuth2Credential getDefaultToken() {
        return poolByUserId.values().stream()
            .filter(pool -> pool.getDefaultToken() != null)
            .min(Comparator.comparingInt(SubscriptionConnectionPool::numSubscriptions))
            .map(IEventSubSocket::getDefaultToken)
            .orElse(null);
    }

    @Override
    public long getLatency() {
        long sum = 0;
        int count = 0;
        for (TwitchSingleUserEventSocketPool pool : poolByUserId.values()) {
            int n = pool.numConnections();
            long latency = pool.getLatency();
            if (latency >= 0) {
                sum += latency * n;
                count += n;
            }
        }
        return count > 0 ? sum / count : -1L;
    }

    /**
     * @return the number of open connections held by this pool.
     */
    public int numConnections() {
        int n = 0;
        for (TwitchSingleUserEventSocketPool pool : poolByUserId.values()) {
            n += pool.numConnections();
        }
        return n;
    }

    /**
     * @return the total number of subscriptions held by all connections.
     */
    public int numSubscriptions() {
        return getSubscriptions().size();
    }

    @Nullable
    private String getUserId(OAuth2Credential token) {
        if (StringUtils.isNotEmpty(token.getUserId())) return token.getUserId();
        identityProvider.getAdditionalCredentialInformation(token).ifPresent(token::updateCredential);
        return token.getUserId();
    }
}