@ -39,7 +39,9 @@ when a child has completed a job.
from __future__ import absolute_import
import logging
import sys
import os . path
import pprint
import threading
import zlib
@ -55,7 +57,7 @@ class Error(Exception):
pass
class ContextService ( mitogen . service . Deduplicating Service) :
class ContextService ( mitogen . service . ) :
"""
Used by workers to fetch the single Context instance corresponding to a
connection configuration , creating the matching connection if it does not
@ -71,33 +73,93 @@ class ContextService(mitogen.service.DeduplicatingService):
"""
handle = 500
max_message_size = 1000
max_contexts = 20
def __init__ ( self , * args , * * kwargs ) :
super ( ContextService , self ) . __init__ ( * args , * * kwargs )
#: JoinPoint for context responses.
self . _lock = threading . Lock ( )
self . _response_by_key = { }
self . _waiters_by_key = { }
#: Active context count.
self . _refs_by_context = { }
#: List of contexts in creation order by via= parameter.
self . _lru_by_via = { }
#: (method_name, kwargs) pairs by Conetxt
self . _cfg_by_context = { }
@mitogen.service.expose ( mitogen . service . AllowParents ( ) )
@mitogen.service.arg_spec ( {
' method_name ' : str
' context' : mitogen . core . Context
} )
def connect ( self , method_name , discriminator = None , * * kwargs ) :
def put( self , context ) :
"""
Return a Context referring to an established connection with the given
configuration , establishing a new connection as necessary .
: param str method_name :
The : class : ` mitogen . parent . Router ` connection method to use .
: param discriminator :
Mixed into the key used to select an existing connection , to allow
intentional duplicate connections to be created .
: param dict kwargs :
Keyword arguments passed to ` mitogen . master . Router . [ method_name ] ( ) ` .
Return a reference , making it eligable for recycling once its reference
count reaches zero .
"""
LOG . debug ( ' %r .put( %r ) ' , self , context )
assert self . _refs_by_context [ context ] > 0
self . _refs_by_context [ context ] - = 1
: returns tuple :
Tuple of ` ( context , home_dir ) ` , where :
* ` context ` is the mitogen . master . Context referring to the
target context .
* ` home_dir ` is a cached copy of the remote directory .
def key_from_kwargs ( self , * * kwargs ) :
"""
Generate a deduplication key from the request . The default
implementation returns a string based on a stable representation of the
input dictionary generated by : py : func : ` pprint . pformat ` .
"""
return pprint . pformat ( kwargs )
def _produce_response ( self , key , response ) :
self . _lock . acquire ( )
try :
waiters = self . _waiters_by_key . pop ( key )
count = len ( waiters )
for msg in waiters :
msg . reply ( response )
finally :
self . _lock . release ( )
return count
def _lru ( self , new_context , * * kwargs ) :
via = kwargs . get ( ' via ' )
if via is None :
# We don't have a limit on the number of directly connections.
return
lru = self . _lru_by_via . setdefault ( via , [ ] )
if len ( lru ) < self . max_contexts :
lru . append ( new_context )
return
for context in reversed ( lru ) :
if self . _refs_by_context [ context ] == 0 :
break
else :
LOG . warning ( ' via= %r reached maximum number of interpreters, '
' but they are all marked as in-use. ' , via )
return
LOG . info ( ' %r ._discard_one(): shutting down %r ' , self , context )
context . shutdown ( )
method_name , kwargs = self . _cfg_by_context [ context ]
key = self . key_from_kwargs ( method_name = method_name , * * kwargs )
self . _lock . acquire ( )
try :
del self . _response_by_key [ key ]
del self . _refs_by_context [ context ]
del self . _cfg_by_context [ context ]
lru . remove ( context )
lru . append ( new_context )
finally :
self . _lock . release ( )
def _connect ( self , method_name , * * kwargs ) :
method = getattr ( self . router , method_name , None )
if method is None :
raise Error ( ' no such Router method: %s ' % ( method_name , ) )
try :
context = method ( * * kwargs )
except mitogen . core . StreamError as e :
@ -107,17 +169,75 @@ class ContextService(mitogen.service.DeduplicatingService):
' msg ' : str ( e ) ,
}
if kwargs . get ( ' via ' ) :
self . _lru ( context , method_name = method_name , * * kwargs )
home_dir = context . call ( os . path . expanduser , ' ~ ' )
# We don't need to wait for the result of this. Ideally we'd check its
# return value somewhere, but logs will catch a failure anyway.
context . call_async ( ansible_mitogen . target . start_fork_parent )
self . _cfg_by_context [ context ] = ( method_name , kwargs )
self . _refs_by_context [ context ] = 0
return {
' context ' : context ,
' home_dir ' : home_dir ,
' msg ' : None ,
}
@mitogen.service.expose ( mitogen . service . AllowParents ( ) )
@mitogen.service.arg_spec ( {
' method_name ' : str
} )
def get ( self , msg , * * kwargs ) :
"""
Return a Context referring to an established connection with the given
configuration , establishing a new connection as necessary .
: param str method_name :
The : class : ` mitogen . parent . Router ` connection method to use .
: param dict kwargs :
Keyword arguments passed to ` mitogen . master . Router . [ method_name ] ( ) ` .
: returns tuple :
Tuple of ` ( context , home_dir ) ` , where :
* ` context ` is the mitogen . master . Context referring to the
target context .
* ` home_dir ` is a cached copy of the remote directory .
"""
key = self . key_from_kwargs ( * * kwargs )
self . _lock . acquire ( )
try :
response = self . _response_by_key . get ( key )
if response is not None :
self . _refs_by_context [ response [ ' context ' ] ] + = 1
return response
waiters = self . _waiters_by_key . get ( key )
if waiters is not None :
waiters . append ( msg )
return self . NO_REPLY
self . _waiters_by_key [ key ] = [ msg ]
finally :
self . _lock . release ( )
# I'm the first thread to wait on a result, so I will create the
# connection.
try :
response = self . _connect ( * * kwargs )
count = self . _produce_response ( key , response )
if response [ ' msg ' ] is None :
self . _response_by_key [ key ] = response
self . _refs_by_context [ response [ ' context ' ] ] + = count
except mitogen . core . CallError :
e = sys . exc_info ( ) [ 1 ]
self . _produce_response ( key , e )
except Exception :
e = sys . exc_info ( ) [ 1 ]
self . _produce_response ( key , mitogen . core . CallError ( e ) )
return self . NO_REPLY
class FileService ( mitogen . service . Service ) :
"""
David Wilson
6 years ago