Retrieve: implement/test stopProducing

[tahoe-lafs/tahoe-lafs.git] / src / allmydata / mutable / retrieve.py

diff --git a/src/allmydata/mutable/retrieve.py b/src/allmydata/mutable/retrieve.py
index 100350a0b62b00925bfae12dddd70b56649665e6..9c09abfecf2aaf06f76eaba6309ab723b2a8ab64 100644 (file)
--- a/src/allmydata/mutable/retrieve.py
+++ b/src/allmydata/mutable/retrieve.py
@@ -7,7 +7,7 @@ from twisted.python import failure
 from twisted.internet.interfaces import IPushProducer, IConsumer
 from foolscap.api import eventually, fireEventually
 from allmydata.interfaces import IRetrieveStatus, NotEnoughSharesError, \
 from twisted.internet.interfaces import IPushProducer, IConsumer
 from foolscap.api import eventually, fireEventually
 from allmydata.interfaces import IRetrieveStatus, NotEnoughSharesError, \

-                                 MDMF_VERSION, SDMF_VERSION
+     DownloadStopped, MDMF_VERSION, SDMF_VERSION

 from allmydata.util import hashutil, log, mathutil
 from allmydata.util.dictutil import DictOfSets
 from allmydata import hashtree, codec
 from allmydata.util import hashutil, log, mathutil
 from allmydata.util.dictutil import DictOfSets
 from allmydata import hashtree, codec


@@ -143,6 +143,7 @@ class Retrieve:
         self._status.set_size(datalength)
         self._status.set_encoding(k, N)
         self.readers = {}
         self._status.set_size(datalength)
         self._status.set_encoding(k, N)
         self.readers = {}

+        self._stopped = False

         self._pause_deferred = None
         self._offset = None
         self._read_length = None
         self._pause_deferred = None
         self._offset = None
         self._read_length = None


@@ -196,6 +197,10 @@ class Retrieve:
 
         eventually(p.callback, None)
 
 
         eventually(p.callback, None)
 

+    def stopProducing(self):
+        self._stopped = True
+        self.resumeProducing()
+

 
     def _check_for_paused(self, res):
         """
 
     def _check_for_paused(self, res):
         """


@@ -205,6 +210,8 @@ class Retrieve:
         the Deferred fires immediately. Otherwise, the Deferred fires
         when the downloader is unpaused.
         """
         the Deferred fires immediately. Otherwise, the Deferred fires
         when the downloader is unpaused.
         """

+        if self._stopped:
+            raise DownloadStopped("our Consumer called stopProducing()")

         if self._pause_deferred is not None:
             d = defer.Deferred()
             self._pause_deferred.addCallback(lambda ignored: d.callback(res))
         if self._pause_deferred is not None:
             d = defer.Deferred()
             self._pause_deferred.addCallback(lambda ignored: d.callback(res))