monolith/modules/chankeep.py

from copy import deepcopy
from math import ceil

from twisted.internet.threads import deferToThread

import main
from utils.logging.debug import debug, trace
from utils.logging.log import error, log, warn


def getEnabledRelays(net):
    """
    Get a list of enabled relays for a network.
    :param net: network
    :rtype: list of int
    :return: list of enabled relay numbers
    """
    enabledRelays = [x for x in main.network[net].relays.keys() if main.network[net].relays[x]["enabled"]]
    debug(f"getEnabledRelays() {net}: {enabledRelays}")
    return enabledRelays


def getConnectedRelays(net):
    """
    Get a list of connected relays for a network.
    :param net: network
    :rtype: list of int
    :return: list of relay numbers
    """
    enabledRelays = getEnabledRelays(net)
    connectedRelays = []
    for i in enabledRelays:
        name = net + str(i)
        if name in main.IRCPool.keys():
            if main.IRCPool[name].isconnected:
                connectedRelays.append(i)
    debug(f"getConnectedRelays() {net}: {connectedRelays}")
    return connectedRelays


def getActiveRelays(net):
    """
    Get a list of active relays for a network.
    :param net: network
    :rtype: list of int
    :return: list of relay numbers
    """
    enabledRelays = getEnabledRelays(net)
    activeRelays = []
    for i in enabledRelays:
        name = net + str(i)
        if name in main.IRCPool.keys():
            debug(
                (
                    f"getActiveRelays() {net}: {i} auth:{main.IRCPool[name].authenticated} "
                    f"conn:{main.IRCPool[name].isconnected}"
                )
            )
            if main.IRCPool[name].authenticated and main.IRCPool[name].isconnected:
                activeRelays.append(i)
    debug(f"getActiveRelays() {net}: {activeRelays}")
    return activeRelays


def relayIsActive(net, num):
    """
    Check if a relay is active.
    :param net: network
    :param num: relay number
    :rtype: bool
    :return: True if relay is active, False otherwise
    """
    activeRelays = getActiveRelays(net)
    return num in activeRelays


def allRelaysActive(net):
    """
    Check if all enabled relays are active and authenticated.
    :param net: network
    :rtype: bool
    :return: True if all relays are active and authenticated, False otherwise
    """
    activeRelays = getActiveRelays(net)
    enabledRelays = getEnabledRelays(net)
    relaysActive = len(activeRelays) == len(enabledRelays)
    debug(f"allRelaysActive() {net}: {relaysActive} ({activeRelays}/{enabledRelays})")
    return relaysActive


def getAverageChanlimit(net):
    """
    Get the average channel limit for a network.
    :param net: network
    :rtype: int
    :return: average channel limit
    """
    total = 0
    for i in getActiveRelays(net):
        name = net + str(i)
        if name in main.IRCPool.keys():
            total += main.IRCPool[name].chanlimit
    avg_chanlimit = total / len(getActiveRelays(net))
    debug(f"getAverageChanlimit() {net}: {avg_chanlimit}")
    return avg_chanlimit


def getSumChanlimit(net):
    """
    Get the sum of all channel limits for a network.
    :param net: network
    :rtype: int
    :return: sum of channel limits
    """
    total = 0
    for i in getActiveRelays(net):
        name = net + str(i)
        if name in main.IRCPool.keys():
            total += main.IRCPool[name].chanlimit
    return total


def getChanFree(net):
    """
    Get a dictionary with the free channel spaces for
    each relay, and a channel limit.
    Example return:
     ({1: 99}, 100)
    :param net: network
    :return: ({relay: channel spaces}, channel limit)
    """
    chanfree = {}
    for i in getActiveRelays(net):
        name = net + str(i)
        if name not in main.IRCPool.keys():
            continue
        if not main.IRCPool[name].isconnected:
            continue
        chanfree[i] = main.IRCPool[name].chanlimit - len(main.IRCPool[name].channels)

    return chanfree


def getTotalChans(net):
    """
    Get the total number of channels on all relays for a network.
    :param net: network
    :rtype: int
    :return: total number of channels
    """
    total = 0
    for i in getActiveRelays(net):
        name = net + str(i)
        if name in main.IRCPool.keys():
            total += len(main.IRCPool[name].channels)
    return total


def emptyChanAllocate(net, flist):
    """
    Allocate channels to relays.
    :param net: network
    :param flist: list of channels to allocate
    :param new: list of newly provisioned relays to account for
    :rtype: dict
    :return: dictionary of {relay: list of channels}"""

    # Get the free channel spaces for each relay
    chanfree = getChanFree(net)
    if not chanfree:
        return

    # Pretend the newly provisioned relays are already on the network
    # for i in new:
    #     chanfree[0][i] = chanfree[1]
    allocated = {}

    newlist = list(flist)
    chan_slots_used = getTotalChans(net)
    max_chans = getSumChanlimit(net) - chan_slots_used
    trunc_list = newlist[:max_chans]
    debug(f"emptyChanAllocate() {net}: newlist:{len(newlist)} trunc_list:{len(trunc_list)}")

    for i in chanfree.keys():
        for x in range(chanfree[i]):
            if not len(trunc_list):
                break
            if i in allocated.keys():
                allocated[i].append(trunc_list.pop())
            else:
                allocated[i] = [trunc_list.pop()]
    return allocated


def populateChans(net, clist):
    """
    Populate channels on relays.
    Stores channels to join in a list in main.TempChan[net][num]
    :param net: network
    :param clist: list of channels to join
    :param new: list of newly provisioned relays to account for"""
    # divided = array_split(clist, relay)
    allocated = emptyChanAllocate(net, clist)
    trace(f"populateChans() allocated:{allocated}")
    if not allocated:
        return
    for i in allocated.keys():
        if net in main.TempChan.keys():
            main.TempChan[net][i] = allocated[i]
        else:
            main.TempChan[net] = {i: allocated[i]}
        trace(f"populateChans() TempChan {net}{i}: {allocated[i]}")


def notifyJoin(net):
    """
    Notify relays to join channels.
    They will pull from main.TempChan and remove channels they join.
    :param net: network
    """
    for i in getActiveRelays(net):
        name = net + str(i)
        if name in main.IRCPool.keys():
            trace(f"notifyJoin() {name}")
            main.IRCPool[name].checkChannels()


def minifyChans(net, listinfo, as_list=False):
    """
    Remove channels from listinfo that are already covered by a relay.
    :param net: network
    :param listinfo: list of channels to check
    :type listinfo: list of [channel, num_users]
    :return: list of channels with joined channels removed
    :rtype: list of [channel, num_users]
    """
    # We want to make this reusable for joining a bunch of channels.
    if as_list:
        channel_list = listinfo

    if not allRelaysActive(net):
        error("All relays for %s are not active, cannot minify list" % net)
        return False
    for i in getActiveRelays(net):
        name = net + str(i)
        for x in main.IRCPool[name].channels:
            if as_list:
                for y in channel_list:
                    if y == x:
                        channel_list.remove(y)
            else:
                for y in listinfo:
                    if y[0] == x:
                        listinfo.remove(y)
    if not as_list:
        if not listinfo:
            log("We're on all the channels we want to be on, dropping LIST")
            return False
    if as_list:
        return channel_list
    else:
        return listinfo


def keepChannels(net, listinfo, mean, sigrelay, relay):
    """
    Minify channels, determine whether we can cover all the channels
    on the network, or need to use 'significant' mode.
    Truncate the channel list to available channel spaces.
    Allocate these channels to relays.
    Notify relays that they should pull from TempChan to join.
    :param net: network
    :param listinfo: list of [channel, num_users] lists
    :param mean: mean of channel population
    :param sigrelay: number of relays needed to cover significant channels
    :param relay: number of relays needed to cover all channels
    :param chanlimit: maximum number of channels to allocate to a relay
    """
    listinfo = minifyChans(net, listinfo)
    if not listinfo:
        return
    if relay <= main.config["ChanKeep"]["SigSwitch"]:  # we can cover all of the channels
        coverAll = True
    elif relay > main.config["ChanKeep"]["SigSwitch"]:  # we cannot cover all of the channels
        coverAll = False
    #    if not sigrelay <= main.config["ChanKeep"]["MaxRelay"]:
    #        error("Network %s is too big to cover: %i relays required" % (net, sigrelay))
    #        return
    num_instances = len(getActiveRelays(net))
    debug(f"keepChannels() {net} instances:{num_instances}")
    chan_slots_used = getTotalChans(net)
    debug(f"keepChannels() slots_used:{chan_slots_used}")
    # max_chans = (chanlimit * num_instances) - chan_slots_used
    max_chans = getSumChanlimit(net) - chan_slots_used
    if max_chans < 0:
        max_chans = 0
    debug(f"keepChannels() max_chans:{max_chans}")
    if coverAll:
        # needed = relay - len(getActiveRelays(net))
        # if needed:
        #     debug(f"keepChannels() coverAll asking to provision {needed} relays for {net} relay:{relay}")
        #     newNums = modules.provision.provisionMultipleRelays(net, needed)
        # else:
        #     newNums = []
        flist = [i[0] for i in listinfo]
        chosen = sorted(flist, reverse=True, key=lambda x: x[1])[:max_chans]
        debug(f"keepChannels() {net}: joining {len(chosen)}/{len(flist)} channels")
        trace(f"keepChannels() {net}: joining:{chosen}")
        populateChans(net, chosen)
    else:
        # needed = sigrelay - len(getActiveRelays(net))
        # if needed:
        #     debug(f"keepChannels() NOT coverAll asking to provision {needed} relays for {net} sigrelay:{sigrelay}")
        #     newNums = modules.provision.provisionMultipleRelays(net, needed)
        # else:
        #     newNums = []
        siglist = [i[0] for i in listinfo if int(i[1]) > mean]
        chosen = sorted(siglist, reverse=True, key=lambda x: x[1])[:max_chans]
        debug(f"keepChannels() {net}: joining {len(chosen)}/{len(siglist)} channels")
        trace(f"keepChannels() {net}: joining:{chosen}")
        populateChans(net, chosen)
    notifyJoin(net)


def joinSingle(net, channel):
    """
    Join a channel on a relay.
    Use ECA to determine which relay to join on.
    :param net: network
    :param channel: channel to join
    :return: relay number that joined the channel
    :rtype: int
    """
    if "," in channel:
        channels = channel.split(",")
        channels = minifyChans(net, channels, as_list=True)

    else:
        channels = [channel]
    populateChans(net, channels)
    notifyJoin(net)
    return True


def partSingle(net, channel):
    """
    Iterate over all the relays of net and part channels matching channel.
    :param net: network
    :param channel: channel to part
    :return: list of relays that parted the channel
    :rtype: list of str
    """
    parted = []
    for i in getConnectedRelays(net):
        name = f"{net}{i}"
        if name in main.IRCPool.keys():
            if channel in main.IRCPool[name].channels:
                main.IRCPool[name].part(channel)
                parted.append(str(i))
    return parted


def nukeNetwork(net):
    """
    Remove network records.
    :param net: network"""
    # purgeRecords(net)
    # p = main.g.pipeline()
    main.g.delete("analytics.list." + net)
    # p.delete("list."+net)
    # p.execute()


# def nukeNetwork(net):
#    deferToThread(_nukeNetwork, net)


def _initialList(net, num, listinfo):
    """
    Called when a relay receives a full LIST response.
    Run statistics to determine how many channels are significant.
    This is done by adding all the numbers of users on the channels together,
    then dividing by the number of channels.
    * cumul - cumulative sum of all channel membership
    * siglength - number of significant channels
    * listlength - number of channels in the list
    * sigrelay - number of relays needed to cover siglength
    * relay - number of relays needed to cover all channels
    :param net: network
    :param num: relay number
    :param listinfo: list of [channel, num_users] lists
    :param chanlimit: maximum number of channels the relay can join
    """
    listlength = len(listinfo)
    cumul = 0
    try:
        cumul += sum(int(i[1]) for i in listinfo)
    except TypeError:
        warn("Bad LIST data received from %s - %i" % (net, num))
        return
    mean = round(cumul / listlength, 2)
    siglength = 0
    insiglength = 0
    sigcumul = 0
    insigcumul = 0
    for i in listinfo:
        if int(i[1]) > mean:
            siglength += 1
            sigcumul += int(i[1])
        elif int(i[1]) < mean:
            insiglength += 1
            insigcumul += int(i[1])

    avg_chanlimit = getAverageChanlimit(net)
    sigrelay = ceil(siglength / avg_chanlimit)
    relay = ceil(listlength / avg_chanlimit)

    cur_relays = len(getActiveRelays(net))
    sig_relays_missing = sigrelay - cur_relays
    all_relays_missing = relay - cur_relays

    abase = "analytics.list.%s" % net
    main.g.delete(abase)
    p = main.g.pipeline()

    # See docstring for meanings
    p.hset(abase, "mean", mean)
    p.hset(abase, "total_chans", listlength)
    p.hset(abase, "big_chans", siglength)
    p.hset(abase, "small_chans", insiglength)
    p.hset(abase, "big_chan_perc", round(siglength / listlength * 100, 2))
    p.hset(abase, "small_chan_perc", round(insiglength / listlength * 100, 2))
    p.hset(abase, "total_cumul_mem", cumul)
    p.hset(abase, "big_chan_cumul_mem", sigcumul)
    p.hset(abase, "small_chan_cumul_mem", insigcumul)
    p.hset(abase, "relays_for_all_chans", relay)
    p.hset(abase, "relays_for_big_chans", sigrelay)
    p.hset(abase, "relays_for_small_chans", ceil(insiglength / avg_chanlimit))
    p.hset(abase, "sig_relays_missing", sig_relays_missing)
    p.hset(abase, "all_relays_missing", all_relays_missing)
    debug(
        (
            f"_initialList() net:{net} num:{num} listlength:{listlength} "
            f"mean:{mean} siglength:{siglength} insiglength:{insiglength} "
            f"sigrelay:{sigrelay} relay:{relay} avg_chanlimit:{avg_chanlimit}"
        )
    )

    # Purge existing records before writing
    # purgeRecords(net)
    # for i in listinfo:
    #    p.rpush(netbase+"."+i[0], i[1])
    #    p.rpush(netbase+"."+i[0], i[2])
    #    p.sadd(netbase, i[0])

    p.execute()
    debug("List parsing completed on %s" % net)
    keepChannels(net, listinfo, mean, sigrelay, relay)

    # return (listinfo, mean, sigrelay, relay)


def convert(data):
    """
    Recursively convert a dictionary.
    """
    if isinstance(data, bytes):
        return data.decode("ascii")
    if isinstance(data, dict):
        return dict(map(convert, data.items()))
    if isinstance(data, tuple):
        return map(convert, data)
    if isinstance(data, list):
        return list(map(convert, data))
    return data


def getListInfo(net):
    abase = f"analytics.list.{net}"
    info = main.g.hgetall(abase)
    return convert(info)


def initialList(net, num, listinfo):
    """
    Run _initialList in a thread.
    See above docstring.
    """
    deferToThread(_initialList, net, num, deepcopy(listinfo))