PortAllocator.cs

Go to the documentation of this file.

using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
 
using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
 
using Tgstation.Server.Host.Configuration;
using Tgstation.Server.Host.Core;
using Tgstation.Server.Host.Database;
using Tgstation.Server.Host.Extensions;
using Tgstation.Server.Host.System;
 
namespace Tgstation.Server.Host.Utils
{
    sealed class PortAllocator : IPortAllocator, IDisposable
    {
        readonly IServerPortProvider serverPortProvider;
 
        readonly IDatabaseContextFactory databaseContextFactory;
 
        readonly IPlatformIdentifier platformIdentifier;
 
        readonly ILogger<PortAllocator> logger;
 
        readonly SwarmConfiguration swarmConfiguration;
 
        readonly SemaphoreSlim allocatorLock;
 
        public PortAllocator(
            IServerPortProvider serverPortProvider,
            IDatabaseContextFactory databaseContextFactory,
            IPlatformIdentifier platformIdentifier,
            IOptions<SwarmConfiguration> swarmConfigurationOptions,
            ILogger<PortAllocator> logger)
        {
            this.serverPortProvider = serverPortProvider ?? throw new ArgumentNullException(nameof(serverPortProvider));
            this.databaseContextFactory = databaseContextFactory ?? throw new ArgumentNullException(nameof(databaseContextFactory));
            this.platformIdentifier = platformIdentifier ?? throw new ArgumentNullException(nameof(platformIdentifier));
            swarmConfiguration = swarmConfigurationOptions?.Value ?? throw new ArgumentNullException(nameof(swarmConfigurationOptions));
            this.logger = logger ?? throw new ArgumentNullException(nameof(logger));
 
            allocatorLock = new SemaphoreSlim(1);
        }
 
        public void Dispose() => allocatorLock.Dispose();
 
        public async ValueTask<ushort?> GetAvailablePort(ushort basePort, bool checkOne, CancellationToken cancellationToken)
        {
            ushort? result = null;
            using (await SemaphoreSlimContext.Lock(allocatorLock, cancellationToken))
                await databaseContextFactory.UseContext(
                    async databaseContext => result = await GetAvailablePort(databaseContext, basePort, checkOne, cancellationToken));
            return result;
        }
 
        async ValueTask<ushort?> GetAvailablePort(IDatabaseContext databaseContext, ushort basePort, bool checkOne, CancellationToken cancellationToken)
        {
            logger.LogTrace("Port allocation >= {basePort} requested...", basePort);
            var ddPorts = await databaseContext
                .DreamDaemonSettings
                .AsQueryable()
                .Where(x => x.Instance!.SwarmIdentifer == swarmConfiguration.Identifier)
                .Select(x => new
                {
                    Port = x.Port!.Value,
                    x.InstanceId,
                })
                .ToListAsync(cancellationToken);
 
            var dmPorts = await databaseContext
                .DreamMakerSettings
                .AsQueryable()
                .Where(x => x.Instance!.SwarmIdentifer == swarmConfiguration.Identifier)
                .Select(x => new
                {
                    ApiValidationPort = x.ApiValidationPort!.Value,
                    x.InstanceId,
                })
                .ToListAsync(cancellationToken);
 
            var exceptions = new List<Exception>();
            ushort port = 0;
            try
            {
                for (port = basePort; port < ushort.MaxValue; ++port)
                {
                    if (checkOne && port != basePort)
                        break;
 
                    if (port == serverPortProvider.HttpApiPort)
                    {
                        logger.LogWarning("Cannot allocate port {port} as it is the TGS API port!", port);
                        continue;
                    }
 
                    var reservedGamePortData = ddPorts.Where(data => data.Port == port).ToList();
                    if (reservedGamePortData.Count > 0)
                    {
                        logger.LogWarning(
                            "Cannot allocate port {port} as it in use by the game server of instance(s): {instanceId}!",
                            port,
                            String.Join(
                                ", ",
                                reservedGamePortData.Select(data => data.InstanceId)));
                        continue;
                    }
 
                    var reservedApiValidationPortData = dmPorts.Where(data => data.ApiValidationPort == port).ToList();
                    if (reservedApiValidationPortData.Count > 0)
                    {
                        logger.LogWarning(
                            "Cannot allocate port {port} as it in use by the API validation server of instance(s): {instanceId}!",
                            port,
                            String.Join(
                                ", ",
                                reservedApiValidationPortData.Select(data => data.InstanceId)));
                        continue;
                    }
 
                    try
                    {
                        SocketExtensions.BindTest(platformIdentifier, port, false, true);
                        SocketExtensions.BindTest(platformIdentifier, port, false, false);
                    }
                    catch (Exception ex)
                    {
                        exceptions.Add(ex);
                        continue;
                    }
 
                    logger.LogInformation("Allocated port {port}", port);
                    return port;
                }
 
                logger.LogWarning("Unable to allocate port >= {basePort}!", basePort);
                return null;
            }
            finally
            {
                if (port != basePort)
                {
                    logger.LogDebug(
                        exceptions.Count == 1
                            ? exceptions.First()
                            : new AggregateException(exceptions),
                        "Failed to allocate ports {basePort}-{lastCheckedPort}!",
                        basePort,
                        port - 1);
                }
            }
        }
    }
}