<commit_before># Copyright (c) 2015 <NAME>
#
# See the file license.txt for copying permission.
import logging
import asyncio
from asyncio import futures
from hbmqtt.mqtt.packet import MQTTFixedHeader
from hbmqtt.mqtt import packet_class
from hbmqtt.errors import NoDataException
from hbmqtt.mqtt.packet import PacketType
from hbmqtt.mqtt.connect import ConnectVariableHeader, ConnectPacket, ConnectPayload
from hbmqtt.mqtt.connack import ConnackPacket
from hbmqtt.mqtt.disconnect import DisconnectPacket
from hbmqtt.mqtt.pingreq import PingReqPacket
from hbmqtt.mqtt.pingresp import PingRespPacket
from hbmqtt.mqtt.publish import PublishPacket
from hbmqtt.mqtt.pubrel import PubrelPacket
from hbmqtt.mqtt.puback import PubackPacket
from hbmqtt.mqtt.pubrec import PubrecPacket
from hbmqtt.mqtt.pubcomp import PubcompPacket
from hbmqtt.mqtt.subscribe import SubscribePacket
from hbmqtt.mqtt.suback import SubackPacket
from hbmqtt.mqtt.unsubscribe import UnsubscribePacket
from hbmqtt.mqtt.unsuback import UnsubackPacket
from hbmqtt.session import Session
from transitions import Machine, MachineError

class InFlightMessage:
    states = ['new', 'published', 'acknowledged', 'received', 'released', 'completed']

    def __init__(self, packet, qos):
        self.packet = packet
        self.qos = qos
        self.puback = None
        self.pubrec = None
        self.pubcomp = None
        self._init_states()

    def _init_states(self):
        self.machine = Machine(model=self, states=InFlightMessage.states, initial='new')
        self.machine.add_transition(trigger='publish', source='new', dest='published')
        if self.qos == 0x01:
            self.machine.add_transition(trigger='acknowledge', source='published', dest='acknowledged')
        if self.qos == 0x02:
            self.machine.add_transition(trigger='receive', source='published', dest='received')
            self.machine.add_transition(trigger='release', source='received', dest='released')
            self.machine.add_transition(trigger='complete', source='released', dest='completed')


class ProtocolHandler:
    """
    Class implementing the MQTT communication protocol using asyncio features
    """

    def __init__(self, session: Session, config, loop=None):
        self.logger = logging.getLogger(__name__)
        self.session = session
        self.config = config
        if loop is None:
            self._loop = asyncio.get_event_loop()
        else:
            self._loop = loop
        self._reader_task = None
        self._writer_task = None
        self._reader_ready = asyncio.Event(loop=self._loop)
        self._writer_ready = asyncio.Event(loop=self._loop)

        self._running = False

        self.session.local_address, self.session.local_port = self.session.writer.get_extra_info('sockname')

        self.incoming_queues = dict()
        self.application_messages = asyncio.Queue()
        for p in PacketType:
            self.incoming_queues[p] = asyncio.Queue()
        self.outgoing_queue = asyncio.Queue()
        self._puback_waiters = dict()
        self._pubrec_waiters = dict()
        self._pubrec_waiters = dict()
        self._pubcomp_waiters = dict()
        self.inflight_messages = dict()

    @asyncio.coroutine
    def start(self):
        self._running = True
        self._reader_task = asyncio.async(self._reader_coro(), loop=self._loop)
        self._writer_task = asyncio.async(self._writer_coro(), loop=self._loop)
        yield from asyncio.wait(
            [self._reader_ready.wait(), self._writer_ready.wait()], loop=self._loop)
        self.logger.debug("Handler tasks started")

    @asyncio.coroutine
    def mqtt_publish(self, topic, message, packet_id, dup, qos, retain):
        if packet_id in self.inflight_messages:
            self.logger.warn("A message with the same packet ID is already in flight")
        packet = PublishPacket.build(topic, message, packet_id, dup, qos, retain)
        yield from self.outgoing_queue.put(packet)
        inflight_message = InFlightMessage(packet, qos)
        self.inflight_messages[packet.variable_header.packet_id] = inflight_message
        inflight_message.publish()
        if qos == 0x01:
            waiter = futures.Future(loop=self._loop)
            self._puback_waiters[packet_id] = waiter
            yield from waiter
            inflight_message.puback = waiter.result()
            inflight_message.acknowledge()
            del self._puback_waiters[packet_id]
        if qos == 0x02:
            # Wait for PUBREC
            waiter = futures.Future(loop=self._loop)
            self._pubrec_waiters[packet_id] = waiter
            yield from waiter
            inflight_message.pubrec = waiter.result()
            del self._pubrec_waiters[packet_id]
            inflight_message.receive()

            # Send pubrel
            pubrel = PubrelPacket.build(packet_id)
            yield from self.outgoing_queue.put(pubrel)
            inflight_message.release()

            # Wait for pubcomp
            waiter = futures.Future(loop=self._loop)
            self._pubcomp_waiters[packet_id] = waiter
            yield from waiter
            inflight_message.pubcomp = waiter.result()
            del self._pubcomp_waiters[packet_id]
            inflight_message.complete()

        del self.inflight_messages[packet_id]
        return inflight_message

    @asyncio.coroutine
    def stop(self):
        self._running = False
        self.session.reader.feed_eof()
        yield from asyncio.wait([self._writer_task, self._reader_task], loop=self._loop)

    @asyncio.coroutine
    def _reader_coro(self):
        self.logger.debug("Starting reader coro")
        while self._running:
            try:
                self._reader_ready.set()
                fixed_header = yield from asyncio.wait_for(MQTTFixedHeader.from_stream(self.session.reader), 5)
                if fixed_header:
                    cls = packet_class(fixed_header)
                    packet = yield from cls.from_stream(self.session.reader, fixed_header=fixed_header)
                    self.logger.debug(" <-in-- " + repr(packet))

                    if packet.fixed_header.packet_type == PacketType.CONNACK:
                        yield from self.handle_connack(packet)
                    if packet.fixed_header.packet_type == PacketType.SUBACK:
                        yield from self.handle_suback(packet)
                    if packet.fixed_header.packet_type == PacketType.UNSUBACK:
                        yield from self.handle_unsuback(packet)
                    if packet.fixed_header.packet_type == PacketType.PUBACK:
                        yield from self.handle_puback(packet)
                    if packet.fixed_header.packet_type == PacketType.PUBREC:
                        yield from self.handle_pubrec(packet)
                    if packet.fixed_header.packet_type == PacketType.PUBCOMP:
                        yield from self.handle_pubcomp(packet)
                    if packet.fixed_header.packet_type == PacketType.PINGRESP:<commit_msg>use Tasks instead of yield from<commit_after># Copyright (c) 2015 <NAME>
#
# See the file license.txt for copying permission.
import logging
import asyncio
from asyncio import futures
from hbmqtt.mqtt.packet import MQTTFixedHeader
from hbmqtt.mqtt import packet_class
from hbmqtt.errors import NoDataException
from hbmqtt.mqtt.packet import PacketType
from hbmqtt.mqtt.connect import ConnectVariableHeader, ConnectPacket, ConnectPayload
from hbmqtt.mqtt.connack import ConnackPacket
from hbmqtt.mqtt.disconnect import DisconnectPacket
from hbmqtt.mqtt.pingreq import PingReqPacket
from hbmqtt.mqtt.pingresp import PingRespPacket
from hbmqtt.mqtt.publish import PublishPacket
from hbmqtt.mqtt.pubrel import PubrelPacket
from hbmqtt.mqtt.puback import PubackPacket
from hbmqtt.mqtt.pubrec import PubrecPacket
from hbmqtt.mqtt.pubcomp import PubcompPacket
from hbmqtt.mqtt.subscribe import SubscribePacket
from hbmqtt.mqtt.suback import SubackPacket
from hbmqtt.mqtt.unsubscribe import UnsubscribePacket
from hbmqtt.mqtt.unsuback import UnsubackPacket
from hbmqtt.session import Session
from transitions import Machine, MachineError

class InFlightMessage:
    states = ['new', 'published', 'acknowledged', 'received', 'released', 'completed']

    def __init__(self, packet, qos):
        self.packet = packet
        self.qos = qos
        self.puback = None
        self.pubrec = None
        self.pubcomp = None
        self._init_states()

    def _init_states(self):
        self.machine = Machine(model=self, states=InFlightMessage.states, initial='new')
        self.machine.add_transition(trigger='publish', source='new', dest='published')
        if self.qos == 0x01:
            self.machine.add_transition(trigger='acknowledge', source='published', dest='acknowledged')
        if self.qos == 0x02:
            self.machine.add_transition(trigger='receive', source='published', dest='received')
            self.machine.add_transition(trigger='release', source='received', dest='released')
            self.machine.add_transition(trigger='complete', source='released', dest='completed')


class ProtocolHandler:
    """
    Class implementing the MQTT communication protocol using asyncio features
    """

    def __init__(self, session: Session, config, loop=None):
        self.logger = logging.getLogger(__name__)
        self.session = session
        self.config = config
        if loop is None:
            self._loop = asyncio.get_event_loop()
        else:
            self._loop = loop
        self._reader_task = None
        self._writer_task = None
        self._reader_ready = asyncio.Event(loop=self._loop)
        self._writer_ready = asyncio.Event(loop=self._loop)

        self._running = False

        self.session.local_address, self.session.local_port = self.session.writer.get_extra_info('sockname')

        self.incoming_queues = dict()
        self.application_messages = asyncio.Queue()
        for p in PacketType:
            self.incoming_queues[p] = asyncio.Queue()
        self.outgoing_queue = asyncio.Queue()
        self._puback_waiters = dict()
        self._pubrec_waiters = dict()
        self._pubrec_waiters = dict()
        self._pubcomp_waiters = dict()
        self.inflight_messages = dict()

    @asyncio.coroutine
    def start(self):
        self._running = True
        self._reader_task = asyncio.async(self._reader_coro(), loop=self._loop)
        self._writer_task = asyncio.async(self._writer_coro(), loop=self._loop)
        yield from asyncio.wait(
            [self._reader_ready.wait(), self._writer_ready.wait()], loop=self._loop)
        self.logger.debug("Handler tasks started")

    @asyncio.coroutine
    def mqtt_publish(self, topic, message, packet_id, dup, qos, retain):
        if packet_id in self.inflight_messages:
            self.logger.warn("A message with the same packet ID is already in flight")
        packet = PublishPacket.build(topic, message, packet_id, dup, qos, retain)
        yield from self.outgoing_queue.put(packet)
        inflight_message = InFlightMessage(packet, qos)
        self.inflight_messages[packet.variable_header.packet_id] = inflight_message
        inflight_message.publish()
        if qos == 0x01:
            waiter = futures.Future(loop=self._loop)
            self._puback_waiters[packet_id] = waiter
            yield from waiter
            inflight_message.puback = waiter.result()
            inflight_message.acknowledge()
            del self._puback_waiters[packet_id]
        if qos == 0x02:
            # Wait for PUBREC
            waiter = futures.Future(loop=self._loop)
            self._pubrec_waiters[packet_id] = waiter
            yield from waiter
            inflight_message.pubrec = waiter.result()
            del self._pubrec_waiters[packet_id]
            inflight_message.receive()

            # Send pubrel
            pubrel = PubrelPacket.build(packet_id)
            yield from self.outgoing_queue.put(pubrel)
            inflight_message.release()

            # Wait for pubcomp
            waiter = futures.Future(loop=self._loop)
            self._pubcomp_waiters[packet_id] = waiter
            yield from waiter
            inflight_message.pubcomp = waiter.result()
            del self._pubcomp_waiters[packet_id]
            inflight_message.complete()

        del self.inflight_messages[packet_id]
        return inflight_message

    @asyncio.coroutine
    def stop(self):
        self._running = False
        self.session.reader.feed_eof()
        yield from asyncio.wait([self._writer_task, self._reader_task], loop=self._loop)

    @asyncio.coroutine
    def _reader_coro(self):
        self.logger.debug("Starting reader coro")
        while self._running:
            try:
                self._reader_ready.set()
                fixed_header = yield from asyncio.wait_for(MQTTFixedHeader.from_stream(self.session.reader), 5)
                if fixed_header:
                    cls = packet_class(fixed_header)
                    packet = yield from cls.from_stream(self.session.reader, fixed_header=fixed_header)
                    self.logger.debug(" <-in-- " + repr(packet))

                    if packet.fixed_header.packet_type == PacketType.CONNACK:
                        asyncio.Task(self.handle_connack(packet))
                    if packet.fixed_header.packet_type == PacketType.SUBACK:
                        asyncio.Task(self.handle_suback(packet))
                    if packet.fixed_header.packet_type == PacketType.UNSUBACK:
                        asyncio.Task(self.handle_unsuback(packet))
                    if packet.fixed_header.packet_type == PacketType.PUBACK:
                        asyncio.Task(self.handle_puback(packet))
                    if packet.fixed_header.packet_type == PacketType.PUBREC:
                        asyncio.Task(self.handle_pubrec(packet))
                    if packet.fixed_header.packet_type == PacketType.PUBCOMP:
                        asyncio.Task(self.handle_pubcomp(packet))
                    if packet.fixed_header.packet_type == PacketType.PINGRESP:
{'fraction_non_alphanumeric': 0.08814374646293152, 'fraction_numerical': 0.0024052065647990947, 'mean_word_length': 2.6027013251783893, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5142504', 'n_tokens_mistral': 4145, 'n_tokens_neox': 3968, 'n_words': 828}