producer/consumerlocator.py

#!/usr/bin/env python

import datetime
import communicator
import os

"""
Consumer locator module, that manages the list of consumers.
"""

__author__ = "@dscharnitzky"
__copyright__ = "Copyright 2020, GoldenPogácsa Team"
__module_name__ = "consumerlocator"
__version__text__ = "1"


class ConsumerLocator:

    """
    Manages the list of consumers.
    """

    def __init__(self):
        """
        Initialize class.
        """
        os.environ["KnownConsumer"] = "10.69.42.2"  # TODO remove
        self.consumerlist = [{"Host": os.environ["KnownConsumer"], "State": True, "LastOk": datetime.datetime.now()}]
        self.currentconsumer = self.consumerlist[0]

    def initcommunicator(self,  comm: communicator.Communicator):
        """
        Initialize the reference to the communicator
        :param comm: is the communicator
        """
        self.communicator = comm

    def learnconsumerlist(self):
        """"
        Learns the list of consumers.
        """
        recievedconsumerlist = self.communicator.discoveravailableconsumers()
        if recievedconsumerlist is None:
            return
        for recconsumer in recievedconsumerlist:
            contains = False
            for consumer in self.consumerlist:
                if consumer["Host"] == recconsumer:
                    contains = True

            if not contains:
                self.consumerlist.append({"Host": recconsumer, "State": True, "LastOk": datetime.datetime.now()})

        self.updateconsumerlist()

    def updateconsumerlist(self):
        """
        Updates the consumer list based on their availability.
        """
        removelist = []
        for consumer in self.consumerlist:
            if not self.communicator.checkconsumer(consumer["Host"]):
                consumer["State"] = False
                if datetime.datetime.now() - consumer["LastOk"] > datetime.timedelta(hours=1):
                    removelist.append(consumer)
            else:
                consumer["LastOk"] = datetime.datetime.now()
                consumer["State"] = True
        for rem in removelist:
            self.consumerlist.remove(rem)

    def updateconsumer(self):
        """
        If the current consumer is not available, checks all the consumers in the list and updates the current one.
        :return: the current consumer or None if there are no available customers at the moment.
        """

        if not self.checkcurrentconsumer():
            self.updateconsumerlist()
            newcurrentconsumer = None

            for consumer in self.consumerlist:
                if consumer["State"]:
                    newcurrentconsumer = consumer
                    break

            self.currentconsumer = newcurrentconsumer
            if self.currentconsumer is not None:
                self.learnconsumerlist()

        if self.currentconsumer is not None:
            return self.currentconsumer["Host"]
        else:
            return None

    def getcurrentconsumer(self):
        """
        Returns the currently selected consumer.
        :return: the current consumer
        """
        return self.currentconsumer["Host"]

    def checkcurrentconsumer(self):
        """
        Check the consumers health.
        :return: True if OK, False if fail
        """
        if self.currentconsumer is None:
            return False
        return self.communicator.checkconsumer(self.currentconsumer["Host"])