|
|
|
|
@ -26,9 +26,15 @@
|
|
|
|
|
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
|
|
|
|
|
# POSSIBILITY OF SUCH DAMAGE.
|
|
|
|
|
|
|
|
|
|
import grp
|
|
|
|
|
import os
|
|
|
|
|
import os.path
|
|
|
|
|
import pprint
|
|
|
|
|
import pwd
|
|
|
|
|
import stat
|
|
|
|
|
import sys
|
|
|
|
|
import threading
|
|
|
|
|
import time
|
|
|
|
|
|
|
|
|
|
import mitogen.core
|
|
|
|
|
import mitogen.select
|
|
|
|
|
@ -37,34 +43,23 @@ from mitogen.core import LOG
|
|
|
|
|
|
|
|
|
|
DEFAULT_POOL_SIZE = 16
|
|
|
|
|
_pool = None
|
|
|
|
|
_pool_pid = None
|
|
|
|
|
#: Serialize pool construction.
|
|
|
|
|
_pool_lock = threading.Lock()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@mitogen.core.takes_router
|
|
|
|
|
def get_or_create_pool(size=None, router=None):
|
|
|
|
|
global _pool
|
|
|
|
|
if _pool is None:
|
|
|
|
|
_pool = Pool(router, [], size=size or DEFAULT_POOL_SIZE)
|
|
|
|
|
return _pool
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@mitogen.core.takes_router
|
|
|
|
|
def _on_stub_call(router):
|
|
|
|
|
"""
|
|
|
|
|
Called for each message received by the core.py stub CALL_SERVICE handler.
|
|
|
|
|
Create the pool if it doesn't already exist, and push enqueued messages
|
|
|
|
|
into the pool's receiver. This may be called more than once as the stub
|
|
|
|
|
service handler runs in asynchronous context, while _on_stub_call() happens
|
|
|
|
|
on the main thread. Multiple CALL_SERVICE may end up enqueued before Pool
|
|
|
|
|
has a chance to install the real CALL_SERVICE handler.
|
|
|
|
|
"""
|
|
|
|
|
pool = get_or_create_pool(router=router)
|
|
|
|
|
mitogen.core._service_call_lock.acquire()
|
|
|
|
|
global _pool_pid
|
|
|
|
|
_pool_lock.acquire()
|
|
|
|
|
try:
|
|
|
|
|
for msg in mitogen.core._service_calls:
|
|
|
|
|
pool._receiver._on_receive(msg)
|
|
|
|
|
del mitogen.core._service_calls[:]
|
|
|
|
|
if _pool_pid != os.getpid():
|
|
|
|
|
_pool = Pool(router, [], size=size or DEFAULT_POOL_SIZE)
|
|
|
|
|
_pool_pid = os.getpid()
|
|
|
|
|
return _pool
|
|
|
|
|
finally:
|
|
|
|
|
mitogen.core._service_call_lock.release()
|
|
|
|
|
_pool_lock.release()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def validate_arg_spec(spec, args):
|
|
|
|
|
@ -149,6 +144,7 @@ class Error(Exception):
|
|
|
|
|
"""
|
|
|
|
|
Raised when an error occurs configuring a service or pool.
|
|
|
|
|
"""
|
|
|
|
|
pass # cope with minify_source() bug.
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class Policy(object):
|
|
|
|
|
@ -183,12 +179,12 @@ class Activator(object):
|
|
|
|
|
|
|
|
|
|
def activate(self, pool, service_name, msg):
|
|
|
|
|
mod_name, _, class_name = service_name.rpartition('.')
|
|
|
|
|
if not self.is_permitted(mod_name, class_name, msg):
|
|
|
|
|
if msg and not self.is_permitted(mod_name, class_name, msg):
|
|
|
|
|
raise mitogen.core.CallError(self.not_active_msg, service_name)
|
|
|
|
|
|
|
|
|
|
module = mitogen.core.import_module(mod_name)
|
|
|
|
|
klass = getattr(module, class_name)
|
|
|
|
|
service = klass(pool.router)
|
|
|
|
|
service = klass(router=pool.router)
|
|
|
|
|
pool.add(service)
|
|
|
|
|
return service
|
|
|
|
|
|
|
|
|
|
@ -238,7 +234,8 @@ class Invoker(object):
|
|
|
|
|
except Exception:
|
|
|
|
|
if no_reply:
|
|
|
|
|
LOG.exception('While calling no-reply method %s.%s',
|
|
|
|
|
type(self).__name__, method.func_name)
|
|
|
|
|
type(self.service).__name__,
|
|
|
|
|
method.func_name)
|
|
|
|
|
else:
|
|
|
|
|
raise
|
|
|
|
|
|
|
|
|
|
@ -249,6 +246,50 @@ class Invoker(object):
|
|
|
|
|
msg.reply(response)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class SerializedInvoker(Invoker):
|
|
|
|
|
def __init__(self, **kwargs):
|
|
|
|
|
super(SerializedInvoker, self).__init__(**kwargs)
|
|
|
|
|
self._lock = threading.Lock()
|
|
|
|
|
self._queue = []
|
|
|
|
|
self._running = False
|
|
|
|
|
|
|
|
|
|
def _pop(self):
|
|
|
|
|
self._lock.acquire()
|
|
|
|
|
try:
|
|
|
|
|
try:
|
|
|
|
|
return self._queue.pop(0)
|
|
|
|
|
except IndexError:
|
|
|
|
|
self._running = False
|
|
|
|
|
finally:
|
|
|
|
|
self._lock.release()
|
|
|
|
|
|
|
|
|
|
def _run(self):
|
|
|
|
|
while True:
|
|
|
|
|
tup = self._pop()
|
|
|
|
|
if tup is None:
|
|
|
|
|
return
|
|
|
|
|
method_name, kwargs, msg = tup
|
|
|
|
|
try:
|
|
|
|
|
super(SerializedInvoker, self).invoke(method_name, kwargs, msg)
|
|
|
|
|
except Exception:
|
|
|
|
|
LOG.exception('%r: while invoking %r of %r',
|
|
|
|
|
self, method_name, self.service)
|
|
|
|
|
msg.reply(mitogen.core.Message.dead())
|
|
|
|
|
|
|
|
|
|
def invoke(self, method_name, kwargs, msg):
|
|
|
|
|
self._lock.acquire()
|
|
|
|
|
try:
|
|
|
|
|
self._queue.append((method_name, kwargs, msg))
|
|
|
|
|
first = not self._running
|
|
|
|
|
self._running = True
|
|
|
|
|
finally:
|
|
|
|
|
self._lock.release()
|
|
|
|
|
|
|
|
|
|
if first:
|
|
|
|
|
self._run()
|
|
|
|
|
return Service.NO_REPLY
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class DeduplicatingInvoker(Invoker):
|
|
|
|
|
"""
|
|
|
|
|
A service that deduplicates and caches expensive responses. Requests are
|
|
|
|
|
@ -398,6 +439,8 @@ class Pool(object):
|
|
|
|
|
thread.start()
|
|
|
|
|
self._threads.append(thread)
|
|
|
|
|
|
|
|
|
|
LOG.debug('%r: initialized', self)
|
|
|
|
|
|
|
|
|
|
@property
|
|
|
|
|
def size(self):
|
|
|
|
|
return len(self._threads)
|
|
|
|
|
@ -407,7 +450,7 @@ class Pool(object):
|
|
|
|
|
if name in self._invoker_by_name:
|
|
|
|
|
raise Error('service named %r already registered' % (name,))
|
|
|
|
|
assert service.select not in self._func_by_recv
|
|
|
|
|
invoker = service.invoker_class(service)
|
|
|
|
|
invoker = service.invoker_class(service=service)
|
|
|
|
|
self._invoker_by_name[name] = invoker
|
|
|
|
|
self._func_by_recv[service.select] = service.on_message
|
|
|
|
|
|
|
|
|
|
@ -427,13 +470,17 @@ class Pool(object):
|
|
|
|
|
invoker = self._invoker_by_name.get(name)
|
|
|
|
|
if not invoker:
|
|
|
|
|
service = self._activator.activate(self, name, msg)
|
|
|
|
|
invoker = service.invoker_class(service)
|
|
|
|
|
invoker = service.invoker_class(service=service)
|
|
|
|
|
self._invoker_by_name[name] = invoker
|
|
|
|
|
finally:
|
|
|
|
|
self._lock.release()
|
|
|
|
|
|
|
|
|
|
return invoker
|
|
|
|
|
|
|
|
|
|
def get_service(self, name):
|
|
|
|
|
invoker = self.get_invoker(name, None)
|
|
|
|
|
return invoker.service
|
|
|
|
|
|
|
|
|
|
def _validate(self, msg):
|
|
|
|
|
tup = msg.unpickle(throw=False)
|
|
|
|
|
if not (isinstance(tup, tuple) and
|
|
|
|
|
@ -454,7 +501,8 @@ class Pool(object):
|
|
|
|
|
LOG.warning('%r: call error: %s: %s', self, msg, e)
|
|
|
|
|
msg.reply(e)
|
|
|
|
|
except Exception:
|
|
|
|
|
LOG.exception('While invoking %r._invoke()', self)
|
|
|
|
|
LOG.exception('%r: while invoking %r of %r',
|
|
|
|
|
self, method_name, service_name)
|
|
|
|
|
e = sys.exc_info()[1]
|
|
|
|
|
msg.reply(mitogen.core.CallError(e))
|
|
|
|
|
|
|
|
|
|
@ -488,3 +536,405 @@ class Pool(object):
|
|
|
|
|
len(self._threads),
|
|
|
|
|
th.name,
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class FileStreamState(object):
|
|
|
|
|
def __init__(self):
|
|
|
|
|
#: List of [(Sender, file object)]
|
|
|
|
|
self.jobs = []
|
|
|
|
|
self.completing = {}
|
|
|
|
|
#: In-flight byte count.
|
|
|
|
|
self.unacked = 0
|
|
|
|
|
#: Lock.
|
|
|
|
|
self.lock = threading.Lock()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class PushFileService(Service):
|
|
|
|
|
"""
|
|
|
|
|
Push-based file service. Files are delivered and cached in RAM, sent
|
|
|
|
|
recursively from parent to child. A child that requests a file via
|
|
|
|
|
:meth:`get` will block until it has ben delivered by a parent.
|
|
|
|
|
|
|
|
|
|
This service will eventually be merged into FileService.
|
|
|
|
|
"""
|
|
|
|
|
invoker_class = SerializedInvoker
|
|
|
|
|
|
|
|
|
|
def __init__(self, **kwargs):
|
|
|
|
|
super(PushFileService, self).__init__(**kwargs)
|
|
|
|
|
self._lock = threading.Lock()
|
|
|
|
|
self._cache = {}
|
|
|
|
|
self._waiters = {}
|
|
|
|
|
self._sent_by_stream = {}
|
|
|
|
|
|
|
|
|
|
def get(self, path):
|
|
|
|
|
self._lock.acquire()
|
|
|
|
|
try:
|
|
|
|
|
if path in self._cache:
|
|
|
|
|
return self._cache[path]
|
|
|
|
|
waiters = self._waiters.setdefault(path, [])
|
|
|
|
|
latch = mitogen.core.Latch()
|
|
|
|
|
waiters.append(lambda: latch.put(None))
|
|
|
|
|
finally:
|
|
|
|
|
self._lock.release()
|
|
|
|
|
|
|
|
|
|
LOG.debug('%r.get(%r) waiting for uncached file to arrive', self, path)
|
|
|
|
|
latch.get()
|
|
|
|
|
LOG.debug('%r.get(%r) -> %r', self, path, self._cache[path])
|
|
|
|
|
return self._cache[path]
|
|
|
|
|
|
|
|
|
|
def _forward(self, context, path):
|
|
|
|
|
stream = self.router.stream_by_id(context.context_id)
|
|
|
|
|
child = mitogen.core.Context(self.router, stream.remote_id)
|
|
|
|
|
sent = self._sent_by_stream.setdefault(stream, set())
|
|
|
|
|
if path in sent and child.context_id != context.context_id:
|
|
|
|
|
child.call_service_async(
|
|
|
|
|
service_name=self.name(),
|
|
|
|
|
method_name='forward',
|
|
|
|
|
path=path,
|
|
|
|
|
context=context
|
|
|
|
|
).close()
|
|
|
|
|
else:
|
|
|
|
|
child.call_service_async(
|
|
|
|
|
service_name=self.name(),
|
|
|
|
|
method_name='store_and_forward',
|
|
|
|
|
path=path,
|
|
|
|
|
data=self._cache[path],
|
|
|
|
|
context=context
|
|
|
|
|
).close()
|
|
|
|
|
|
|
|
|
|
@expose(policy=AllowParents())
|
|
|
|
|
@arg_spec({
|
|
|
|
|
'context': mitogen.core.Context,
|
|
|
|
|
'paths': list,
|
|
|
|
|
'modules': list,
|
|
|
|
|
})
|
|
|
|
|
def propagate_paths_and_modules(self, context, paths, modules):
|
|
|
|
|
"""
|
|
|
|
|
One size fits all method to ensure a target context has been preloaded
|
|
|
|
|
with a set of small files and Python modules.
|
|
|
|
|
"""
|
|
|
|
|
for path in paths:
|
|
|
|
|
self.propagate_to(context, path)
|
|
|
|
|
for fullname in modules:
|
|
|
|
|
self.router.responder.forward_module(context, fullname)
|
|
|
|
|
|
|
|
|
|
@expose(policy=AllowParents())
|
|
|
|
|
@arg_spec({
|
|
|
|
|
'context': mitogen.core.Context,
|
|
|
|
|
'path': basestring,
|
|
|
|
|
})
|
|
|
|
|
def propagate_to(self, context, path):
|
|
|
|
|
LOG.debug('%r.propagate_to(%r, %r)', self, context, path)
|
|
|
|
|
if path not in self._cache:
|
|
|
|
|
fp = open(path, 'rb')
|
|
|
|
|
try:
|
|
|
|
|
self._cache[path] = mitogen.core.Blob(fp.read())
|
|
|
|
|
finally:
|
|
|
|
|
fp.close()
|
|
|
|
|
self._forward(context, path)
|
|
|
|
|
|
|
|
|
|
def _store(self, path, data):
|
|
|
|
|
self._lock.acquire()
|
|
|
|
|
try:
|
|
|
|
|
self._cache[path] = data
|
|
|
|
|
return self._waiters.pop(path, [])
|
|
|
|
|
finally:
|
|
|
|
|
self._lock.release()
|
|
|
|
|
|
|
|
|
|
@expose(policy=AllowParents())
|
|
|
|
|
@no_reply()
|
|
|
|
|
@arg_spec({
|
|
|
|
|
'path': basestring,
|
|
|
|
|
'data': mitogen.core.Blob,
|
|
|
|
|
'context': mitogen.core.Context,
|
|
|
|
|
})
|
|
|
|
|
def store_and_forward(self, path, data, context):
|
|
|
|
|
LOG.debug('%r.store_and_forward(%r, %r, %r)',
|
|
|
|
|
self, path, data, context)
|
|
|
|
|
waiters = self._store(path, data)
|
|
|
|
|
if context.context_id != mitogen.context_id:
|
|
|
|
|
self._forward(context, path)
|
|
|
|
|
for callback in waiters:
|
|
|
|
|
callback()
|
|
|
|
|
|
|
|
|
|
@expose(policy=AllowParents())
|
|
|
|
|
@no_reply()
|
|
|
|
|
@arg_spec({
|
|
|
|
|
'path': basestring,
|
|
|
|
|
'context': mitogen.core.Context,
|
|
|
|
|
})
|
|
|
|
|
def forward(self, path, context):
|
|
|
|
|
LOG.debug('%r.forward(%r, %r)', self, path, context)
|
|
|
|
|
if path not in self._cache:
|
|
|
|
|
LOG.error('%r: %r is not in local cache', self, path)
|
|
|
|
|
return
|
|
|
|
|
self._forward(path, context)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class FileService(Service):
|
|
|
|
|
"""
|
|
|
|
|
Streaming file server, used to serve small and huge files alike. Paths must
|
|
|
|
|
be registered by a trusted context before they will be served to a child.
|
|
|
|
|
|
|
|
|
|
Transfers are divided among the physical streams that connect external
|
|
|
|
|
contexts, ensuring each stream never has excessive data buffered in RAM,
|
|
|
|
|
while still maintaining enough to fully utilize available bandwidth. This
|
|
|
|
|
is achieved by making an initial bandwidth assumption, enqueueing enough
|
|
|
|
|
chunks to fill that assumed pipe, then responding to delivery
|
|
|
|
|
acknowledgements from the receiver by scheduling new chunks.
|
|
|
|
|
|
|
|
|
|
Transfers proceed one-at-a-time per stream. When multiple contexts exist on
|
|
|
|
|
a stream (e.g. one is the SSH account, another is a sudo account, and a
|
|
|
|
|
third is a proxied SSH connection), each request is satisfied in turn
|
|
|
|
|
before subsequent requests start flowing. This ensures when a stream is
|
|
|
|
|
contended, priority is given to completing individual transfers rather than
|
|
|
|
|
potentially aborting many partial transfers, causing the bandwidth to be
|
|
|
|
|
wasted.
|
|
|
|
|
|
|
|
|
|
Theory of operation:
|
|
|
|
|
1. Trusted context (i.e. WorkerProcess) calls register(), making a
|
|
|
|
|
file available to any untrusted context.
|
|
|
|
|
2. Requestee context creates a mitogen.core.Receiver() to receive
|
|
|
|
|
chunks, then calls fetch(path, recv.to_sender()), to set up the
|
|
|
|
|
transfer.
|
|
|
|
|
3. fetch() replies to the call with the file's metadata, then
|
|
|
|
|
schedules an initial burst up to the window size limit (1MiB).
|
|
|
|
|
4. Chunks begin to arrive in the requestee, which calls acknowledge()
|
|
|
|
|
for each 128KiB received.
|
|
|
|
|
5. The acknowledge() call arrives at FileService, which scheduled a new
|
|
|
|
|
chunk to refill the drained window back to the size limit.
|
|
|
|
|
6. When the last chunk has been pumped for a single transfer,
|
|
|
|
|
Sender.close() is called causing the receive loop in
|
|
|
|
|
target.py::_get_file() to exit, allowing that code to compare the
|
|
|
|
|
transferred size with the total file size from the metadata.
|
|
|
|
|
7. If the sizes mismatch, _get_file()'s caller is informed which will
|
|
|
|
|
discard the result and log/raise an error.
|
|
|
|
|
|
|
|
|
|
Shutdown:
|
|
|
|
|
1. process.py calls service.Pool.shutdown(), which arranges for the
|
|
|
|
|
service pool threads to exit and be joined, guranteeing no new
|
|
|
|
|
requests can arrive, before calling Service.on_shutdown() for each
|
|
|
|
|
registered service.
|
|
|
|
|
2. FileService.on_shutdown() walks every in-progress transfer and calls
|
|
|
|
|
Sender.close(), causing Receiver loops in the requestees to exit
|
|
|
|
|
early. The size check fails and any partially downloaded file is
|
|
|
|
|
discarded.
|
|
|
|
|
3. Control exits _get_file() in every target, and graceful shutdown can
|
|
|
|
|
proceed normally, without the associated thread needing to be
|
|
|
|
|
forcefully killed.
|
|
|
|
|
"""
|
|
|
|
|
unregistered_msg = 'Path is not registered with FileService.'
|
|
|
|
|
context_mismatch_msg = 'sender= kwarg context must match requestee context'
|
|
|
|
|
|
|
|
|
|
#: Burst size. With 1MiB and 10ms RTT max throughput is 100MiB/sec, which
|
|
|
|
|
#: is 5x what SSH can handle on a 2011 era 2.4Ghz Core i5.
|
|
|
|
|
window_size_bytes = 1048576
|
|
|
|
|
|
|
|
|
|
def __init__(self, router):
|
|
|
|
|
super(FileService, self).__init__(router)
|
|
|
|
|
#: Mapping of registered path -> file size.
|
|
|
|
|
self._metadata_by_path = {}
|
|
|
|
|
#: Mapping of Stream->FileStreamState.
|
|
|
|
|
self._state_by_stream = {}
|
|
|
|
|
|
|
|
|
|
def _name_or_none(self, func, n, attr):
|
|
|
|
|
try:
|
|
|
|
|
return getattr(func(n), attr)
|
|
|
|
|
except KeyError:
|
|
|
|
|
return None
|
|
|
|
|
|
|
|
|
|
@expose(policy=AllowParents())
|
|
|
|
|
@arg_spec({
|
|
|
|
|
'path': basestring,
|
|
|
|
|
})
|
|
|
|
|
def register(self, path):
|
|
|
|
|
"""
|
|
|
|
|
Authorize a path for access by children. Repeat calls with the same
|
|
|
|
|
path is harmless.
|
|
|
|
|
|
|
|
|
|
:param str path:
|
|
|
|
|
File path.
|
|
|
|
|
"""
|
|
|
|
|
if path in self._metadata_by_path:
|
|
|
|
|
return
|
|
|
|
|
|
|
|
|
|
st = os.stat(path)
|
|
|
|
|
if not stat.S_ISREG(st.st_mode):
|
|
|
|
|
raise IOError('%r is not a regular file.' % (in_path,))
|
|
|
|
|
|
|
|
|
|
LOG.debug('%r: registering %r', self, path)
|
|
|
|
|
self._metadata_by_path[path] = {
|
|
|
|
|
'size': st.st_size,
|
|
|
|
|
'mode': st.st_mode,
|
|
|
|
|
'owner': self._name_or_none(pwd.getpwuid, 0, 'pw_name'),
|
|
|
|
|
'group': self._name_or_none(grp.getgrgid, 0, 'gr_name'),
|
|
|
|
|
'mtime': st.st_mtime,
|
|
|
|
|
'atime': st.st_atime,
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
def on_shutdown(self):
|
|
|
|
|
"""
|
|
|
|
|
Respond to shutdown by sending close() to every target, allowing their
|
|
|
|
|
receive loop to exit and clean up gracefully.
|
|
|
|
|
"""
|
|
|
|
|
LOG.debug('%r.on_shutdown()', self)
|
|
|
|
|
for stream, state in self._state_by_stream.items():
|
|
|
|
|
state.lock.acquire()
|
|
|
|
|
try:
|
|
|
|
|
for sender, fp in reversed(state.jobs):
|
|
|
|
|
sender.close()
|
|
|
|
|
fp.close()
|
|
|
|
|
state.jobs.pop()
|
|
|
|
|
finally:
|
|
|
|
|
state.lock.release()
|
|
|
|
|
|
|
|
|
|
# The IO loop pumps 128KiB chunks. An ideal message is a multiple of this,
|
|
|
|
|
# odd-sized messages waste one tiny write() per message on the trailer.
|
|
|
|
|
# Therefore subtract 10 bytes pickle overhead + 24 bytes header.
|
|
|
|
|
IO_SIZE = mitogen.core.CHUNK_SIZE - (mitogen.core.Stream.HEADER_LEN + (
|
|
|
|
|
len(
|
|
|
|
|
mitogen.core.Message.pickled(
|
|
|
|
|
mitogen.core.Blob(' ' * mitogen.core.CHUNK_SIZE)
|
|
|
|
|
).data
|
|
|
|
|
) - mitogen.core.CHUNK_SIZE
|
|
|
|
|
))
|
|
|
|
|
|
|
|
|
|
def _schedule_pending_unlocked(self, state):
|
|
|
|
|
"""
|
|
|
|
|
Consider the pending transfers for a stream, pumping new chunks while
|
|
|
|
|
the unacknowledged byte count is below :attr:`window_size_bytes`. Must
|
|
|
|
|
be called with the FileStreamState lock held.
|
|
|
|
|
|
|
|
|
|
:param FileStreamState state:
|
|
|
|
|
Stream to schedule chunks for.
|
|
|
|
|
"""
|
|
|
|
|
while state.jobs and state.unacked < self.window_size_bytes:
|
|
|
|
|
sender, fp = state.jobs[0]
|
|
|
|
|
s = fp.read(self.IO_SIZE)
|
|
|
|
|
if s:
|
|
|
|
|
state.unacked += len(s)
|
|
|
|
|
sender.send(mitogen.core.Blob(s))
|
|
|
|
|
else:
|
|
|
|
|
# File is done. Cause the target's receive loop to exit by
|
|
|
|
|
# closing the sender, close the file, and remove the job entry.
|
|
|
|
|
sender.close()
|
|
|
|
|
fp.close()
|
|
|
|
|
state.jobs.pop(0)
|
|
|
|
|
|
|
|
|
|
@expose(policy=AllowAny())
|
|
|
|
|
@no_reply()
|
|
|
|
|
@arg_spec({
|
|
|
|
|
'path': basestring,
|
|
|
|
|
'sender': mitogen.core.Sender,
|
|
|
|
|
})
|
|
|
|
|
def fetch(self, path, sender, msg):
|
|
|
|
|
"""
|
|
|
|
|
Start a transfer for a registered path.
|
|
|
|
|
|
|
|
|
|
:param str path:
|
|
|
|
|
File path.
|
|
|
|
|
:param mitogen.core.Sender sender:
|
|
|
|
|
Sender to receive file data.
|
|
|
|
|
:returns:
|
|
|
|
|
Dict containing the file metadata:
|
|
|
|
|
|
|
|
|
|
* ``size``: File size in bytes.
|
|
|
|
|
* ``mode``: Integer file mode.
|
|
|
|
|
* ``owner``: Owner account name on host machine.
|
|
|
|
|
* ``group``: Owner group name on host machine.
|
|
|
|
|
* ``mtime``: Floating point modification time.
|
|
|
|
|
* ``ctime``: Floating point change time.
|
|
|
|
|
:raises Error:
|
|
|
|
|
Unregistered path, or Sender did not match requestee context.
|
|
|
|
|
"""
|
|
|
|
|
if path not in self._metadata_by_path:
|
|
|
|
|
raise Error(self.unregistered_msg)
|
|
|
|
|
if msg.src_id != sender.context.context_id:
|
|
|
|
|
raise Error(self.context_mismatch_msg)
|
|
|
|
|
|
|
|
|
|
LOG.debug('Serving %r', path)
|
|
|
|
|
fp = open(path, 'rb', self.IO_SIZE)
|
|
|
|
|
# Response must arrive first so requestee can begin receive loop,
|
|
|
|
|
# otherwise first ack won't arrive until all pending chunks were
|
|
|
|
|
# delivered. In that case max BDP would always be 128KiB, aka. max
|
|
|
|
|
# ~10Mbit/sec over a 100ms link.
|
|
|
|
|
msg.reply(self._metadata_by_path[path])
|
|
|
|
|
|
|
|
|
|
stream = self.router.stream_by_id(sender.context.context_id)
|
|
|
|
|
state = self._state_by_stream.setdefault(stream, FileStreamState())
|
|
|
|
|
state.lock.acquire()
|
|
|
|
|
try:
|
|
|
|
|
state.jobs.append((sender, fp))
|
|
|
|
|
self._schedule_pending_unlocked(state)
|
|
|
|
|
finally:
|
|
|
|
|
state.lock.release()
|
|
|
|
|
|
|
|
|
|
@expose(policy=AllowAny())
|
|
|
|
|
@no_reply()
|
|
|
|
|
@arg_spec({
|
|
|
|
|
'size': int,
|
|
|
|
|
})
|
|
|
|
|
@no_reply()
|
|
|
|
|
def acknowledge(self, size, msg):
|
|
|
|
|
"""
|
|
|
|
|
Acknowledge bytes received by a transfer target, scheduling new chunks
|
|
|
|
|
to keep the window full. This should be called for every chunk received
|
|
|
|
|
by the target.
|
|
|
|
|
"""
|
|
|
|
|
stream = self.router.stream_by_id(msg.src_id)
|
|
|
|
|
state = self._state_by_stream[stream]
|
|
|
|
|
state.lock.acquire()
|
|
|
|
|
try:
|
|
|
|
|
if state.unacked < size:
|
|
|
|
|
LOG.error('%r.acknowledge(src_id %d): unacked=%d < size %d',
|
|
|
|
|
self, msg.src_id, state.unacked, size)
|
|
|
|
|
state.unacked -= min(state.unacked, size)
|
|
|
|
|
self._schedule_pending_unlocked(state)
|
|
|
|
|
finally:
|
|
|
|
|
state.lock.release()
|
|
|
|
|
|
|
|
|
|
@classmethod
|
|
|
|
|
def get(cls, context, path, out_fp):
|
|
|
|
|
"""
|
|
|
|
|
Streamily download a file from the connection multiplexer process in
|
|
|
|
|
the controller.
|
|
|
|
|
|
|
|
|
|
:param mitogen.core.Context context:
|
|
|
|
|
Reference to the context hosting the FileService that will be used
|
|
|
|
|
to fetch the file.
|
|
|
|
|
:param bytes in_path:
|
|
|
|
|
FileService registered name of the input file.
|
|
|
|
|
:param bytes out_path:
|
|
|
|
|
Name of the output path on the local disk.
|
|
|
|
|
:returns:
|
|
|
|
|
:data:`True` on success, or :data:`False` if the transfer was
|
|
|
|
|
interrupted and the output should be discarded.
|
|
|
|
|
"""
|
|
|
|
|
LOG.debug('get_file(): fetching %r from %r', path, context)
|
|
|
|
|
t0 = time.time()
|
|
|
|
|
recv = mitogen.core.Receiver(router=context.router)
|
|
|
|
|
metadata = context.call_service(
|
|
|
|
|
service_name=cls.name(),
|
|
|
|
|
method_name='fetch',
|
|
|
|
|
path=path,
|
|
|
|
|
sender=recv.to_sender(),
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
for chunk in recv:
|
|
|
|
|
s = chunk.unpickle()
|
|
|
|
|
LOG.debug('get_file(%r): received %d bytes', path, len(s))
|
|
|
|
|
context.call_service_async(
|
|
|
|
|
service_name=cls.name(),
|
|
|
|
|
method_name='acknowledge',
|
|
|
|
|
size=len(s),
|
|
|
|
|
).close()
|
|
|
|
|
out_fp.write(s)
|
|
|
|
|
|
|
|
|
|
ok = out_fp.tell() == metadata['size']
|
|
|
|
|
if not ok:
|
|
|
|
|
LOG.error('get_file(%r): receiver was closed early, controller '
|
|
|
|
|
'is likely shutting down.', path)
|
|
|
|
|
|
|
|
|
|
LOG.debug('target.get_file(): fetched %d bytes of %r from %r in %dms',
|
|
|
|
|
metadata['size'], path, context, 1000 * (time.time() - t0))
|
|
|
|
|
return ok, metadata
|
|
|
|
|
|
dw
8 years ago
committed by