Mercurial > code > home > repos > homeauto
annotate lib/mqtt_client/mqtt_client.py @ 1382:f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
Ignore-this: 301b82746e517d2a6ff212677f23ca8e
darcs-hash:c3badf7258cd931afa0f9b0507482d0c6a702407
author | drewp <drewp@bigasterisk.com> |
---|---|
date | Wed, 08 May 2019 00:55:58 -0700 |
parents | 3cf19717cb6f |
children | c887b1cc5e83 |
rev | line source |
---|---|
1357 | 1 import logging |
2 from mqtt.client.factory import MQTTFactory | |
3 from rx import Observable | |
1382
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
4 from rx.subjects import Subject |
1357 | 5 from rx.concurrency import TwistedScheduler |
1382
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
6 from twisted.application.internet import ClientService |
1357 | 7 from twisted.internet import reactor |
8 from twisted.internet.defer import inlineCallbacks, Deferred | |
9 from twisted.internet.endpoints import clientFromString | |
10 | |
11 log = logging.getLogger('mqtt_client') | |
12 | |
13 class MQTTService(ClientService): | |
14 | |
1382
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
15 def __init__(self, endpoint, factory, observersByTopic): |
1357 | 16 self.endpoint = endpoint |
1382
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
17 self.observersByTopic = observersByTopic |
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
18 ClientService.__init__(self, endpoint, factory, retryPolicy=lambda _: 5) |
1357 | 19 |
20 def startService(self): | |
21 self.whenConnected().addCallback(self.connectToBroker) | |
22 ClientService.startService(self) | |
23 | |
1382
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
24 def ensureSubscribed(self, topic: bytes): |
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
25 self.whenConnected().addCallback(self._subscribeToLatestTopic, topic) |
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
26 |
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
27 def _subscribeToLatestTopic(self, protocol, topic: bytes): |
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
28 if protocol.state == protocol.CONNECTED: |
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
29 self.protocol.subscribe(topics=[(topic.decode('utf8'), 2)]) |
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
30 # else it'll get done in the next connectToBroker. |
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
31 |
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
32 def _subscribeAll(self): |
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
33 topics = list(self.observersByTopic) |
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
34 log.info('subscribing %r', topics) |
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
35 self.protocol.subscribe(topics=[(topic.decode('utf8'), 2) for topic in topics]) |
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
36 |
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
37 |
1357 | 38 @inlineCallbacks |
39 def connectToBroker(self, protocol): | |
40 self.protocol = protocol | |
1382
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
41 self.protocol.onDisconnection = self._onProtocolDisconnection |
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
42 |
1357 | 43 # Publish requests beyond window size are enqueued |
44 self.protocol.setWindowSize(1) | |
45 | |
46 try: | |
47 yield self.protocol.connect("TwistedMQTT-pub", keepalive=60) | |
48 except Exception as e: | |
1382
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
49 log.error(f"Connecting to {self.endpoint} raised {e!s}") |
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
50 return |
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
51 |
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
52 log.info(f"Connected to {self.endpoint}") |
1357 | 53 |
1382
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
54 self.protocol.onPublish = self._onProtocolMessage |
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
55 self._subscribeAll() |
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
56 |
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
57 def _onProtocolMessage(self, topic, payload, qos, dup, retain, msgId): |
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
58 topic = topic.encode('ascii') |
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
59 observers = self.observersByTopic.get(topic, []) |
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
60 log.debug(f'received {topic} payload {payload} ({len(observers)} obs)') |
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
61 for obs in observers: |
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
62 obs.on_next(payload) |
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
63 |
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
64 def _onProtocolDisconnection(self, reason): |
1357 | 65 log.warn("Connection to broker lost: %r", reason) |
66 self.whenConnected().addCallback(self.connectToBroker) | |
67 | |
1359 | 68 def publish(self, topic: bytes, msg: bytes): |
1357 | 69 def _logFailure(failure): |
70 log.warn("publish failed: %s", failure.getErrorMessage()) | |
71 return failure | |
72 | |
1363
3cf19717cb6f
also pass topic as str at publish
drewp <drewp@bigasterisk.com>
parents:
1361
diff
changeset
|
73 return self.protocol.publish(topic=topic.decode('utf-8'), qos=0, |
1359 | 74 message=bytearray(msg)).addErrback(_logFailure) |
1357 | 75 |
76 | |
77 class MqttClient(object): | |
78 def __init__(self, brokerHost='bang', brokerPort=1883): | |
79 | |
1382
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
80 self.observersByTopic = {} # bytes: Set(observer) |
1357 | 81 |
82 factory = MQTTFactory(profile=MQTTFactory.PUBLISHER | MQTTFactory.SUBSCRIBER) | |
83 myEndpoint = clientFromString(reactor, 'tcp:%s:%s' % (brokerHost, brokerPort)) | |
84 myEndpoint.__class__.__repr__ = lambda self: repr('%s:%s' % (self._host, self._port)) | |
1382
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
85 self.serv = MQTTService(myEndpoint, factory, self.observersByTopic) |
1357 | 86 self.serv.startService() |
1382
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
87 |
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
88 def publish(self, topic: bytes, msg: bytes): |
1357 | 89 return self.serv.publish(topic, msg) |
90 | |
1359 | 91 def subscribe(self, topic: bytes): |
1357 | 92 """returns rx.Observable of payload strings""" |
1382
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
93 ret = Subject() |
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
94 self.observersByTopic.setdefault(topic, set()).add(ret) |
f883166f7ca1
big rewrite. now probably works for multiple subscriptions and over reconnects
drewp <drewp@bigasterisk.com>
parents:
1363
diff
changeset
|
95 self.serv.ensureSubscribed(topic) |
1357 | 96 return ret |