summary refs log tree commit diff
path: root/synapse/util/file_consumer.py
blob: d19d48665cd1efc51fc1bf2257be71aaa78b79a9 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
# -*- coding: utf-8 -*-
# Copyright 2018 New Vector Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from twisted.internet import defer, threads, reactor

from synapse.util.logcontext import make_deferred_yieldable

import Queue


class BackgroundFileConsumer(object):
    """A consumer that writes to a file like object. Supports both push
    and pull producers

    Args:
        file_obj (file): The file like object to write to. Closed when
            finished.
    """

    # For PushProducers pause if we have this many unwritten slices
    _PAUSE_ON_QUEUE_SIZE = 5
    # And resume once the size of the queue is less than this
    _RESUME_ON_QUEUE_SIZE = 2

    def __init__(self, file_obj):
        self._file_obj = file_obj

        # Producer we're registered with
        self._producer = None

        # True if PushProducer, false if PullProducer
        self.streaming = False

        # For PushProducers, indicates whether we've paused the producer and
        # need to call resumeProducing before we get more data.
        self._paused_producer = False

        # Queue of slices of bytes to be written. When producer calls
        # unregister a final None is sent.
        self._bytes_queue = Queue.Queue()

        # Deferred that is resolved when finished writing
        self._finished_deferred = None

        # If the _writer thread throws an exception it gets stored here.
        self._write_exception = None

        # A deferred that gets resolved when the bytes_queue gets empty.
        # Mainly used for tests.
        self._notify_empty_deferred = None

    def registerProducer(self, producer, streaming):
        """Part of IConsumer interface

        Args:
            producer (IProducer)
            streaming (bool): True if push based producer, False if pull
                based.
        """
        if self._producer:
            raise Exception("registerProducer called twice")

        self._producer = producer
        self.streaming = streaming
        self._finished_deferred = threads.deferToThread(self._writer)
        if not streaming:
            self._producer.resumeProducing()

    def unregisterProducer(self):
        """Part of IProducer interface
        """
        self._producer = None
        if not self._finished_deferred.called:
            self._bytes_queue.put_nowait(None)

    def write(self, bytes):
        """Part of IProducer interface
        """
        if self._write_exception:
            raise self._write_exception

        if self._finished_deferred.called:
            raise Exception("consumer has closed")

        self._bytes_queue.put_nowait(bytes)

        # If this is a PushProducer and the queue is getting behind
        # then we pause the producer.
        if self.streaming and self._bytes_queue.qsize() >= self._PAUSE_ON_QUEUE_SIZE:
            self._paused_producer = True
            self._producer.pauseProducing()

    def _writer(self):
        """This is run in a background thread to write to the file.
        """
        try:
            while self._producer or not self._bytes_queue.empty():
                # If we've paused the producer check if we should resume the
                # producer.
                if self._producer and self._paused_producer:
                    if self._bytes_queue.qsize() <= self._RESUME_ON_QUEUE_SIZE:
                        reactor.callFromThread(self._resume_paused_producer)

                if self._notify_empty_deferred and self._bytes_queue.empty():
                    reactor.callFromThread(self._notify_empty)

                bytes = self._bytes_queue.get()

                # If we get a None (or empty list) then that's a signal used
                # to indicate we should check if we should stop.
                if bytes:
                    self._file_obj.write(bytes)

                # If its a pull producer then we need to explicitly ask for
                # more stuff.
                if not self.streaming and self._producer:
                    reactor.callFromThread(self._producer.resumeProducing)
        except Exception as e:
            self._write_exception = e
            raise
        finally:
            self._file_obj.close()

    def wait(self):
        """Returns a deferred that resolves when finished writing to file
        """
        return make_deferred_yieldable(self._finished_deferred)

    def _resume_paused_producer(self):
        """Gets called if we should resume producing after being paused
        """
        if self._paused_producer and self._producer:
            self._paused_producer = False
            self._producer.resumeProducing()

    def _notify_empty(self):
        """Called when the _writer thread thinks the queue may be empty and
        we should notify anything waiting on `wait_for_writes`
        """
        if self._notify_empty_deferred and self._bytes_queue.empty():
            d = self._notify_empty_deferred
            self._notify_empty_deferred = None
            d.callback(None)

    def wait_for_writes(self):
        """Wait for the write queue to be empty and for writes to have
        finished. This is mainly useful for tests.
        """
        if not self._notify_empty_deferred:
            self._notify_empty_deferred = defer.Deferred()
        return self._notify_empty_deferred