<s> from collections import deque <EOL> import logging <EOL> import gevent <EOL> from . base import PoolSink <EOL> from . . async import AsyncResult <EOL> from . . constants import ( Int , ChannelState , SinkProperties , SinkRole ) <EOL> from . . sink import ( <EOL> ClientMessageSink , <EOL> SinkProvider , <EOL> FailingMessageSink <EOL> ) <EOL> from . . dispatch import ServiceClosedError <EOL> from . . varz import ( <EOL> Gauge , <EOL> Source , <EOL> VarzBase <EOL> ) <EOL> class QueuingMessageSink ( ClientMessageSink ) : <EOL> def __init__ ( self , queue ) : <EOL> super ( QueuingMessageSink , self ) . __init__ ( ) <EOL> self . _queue = queue <EOL> def AsyncProcessRequest ( self , sink_stack , msg , stream , headers ) : <EOL> self . _queue . append ( ( sink_stack , msg , stream , headers ) ) <EOL> def AsyncProcessResponse ( self , sink_stack , context , stream , msg ) : <EOL> raise NotImplementedError ( "<STR_LIT>" ) <EOL> @ property <EOL> def state ( self ) : <EOL> return ChannelState . Open <EOL> class MaxWaitersError ( Exception ) : pass <EOL> class WatermarkPoolSink ( PoolSink ) : <EOL> """<STR_LIT>""" <EOL> ROOT_LOG = logging . getLogger ( '<STR_LIT>' ) <EOL> class Varz ( VarzBase ) : <EOL> """<STR_LIT>""" <EOL> _VARZ_BASE_NAME = '<STR_LIT>' <EOL> _VARZ = { <EOL> '<STR_LIT:size>' : Gauge , <EOL> '<STR_LIT>' : Gauge <EOL> } <EOL> def __init__ ( self , next_provider , sink_properties , global_properties ) : <EOL> endpoint = global_properties [ SinkProperties . Endpoint ] <EOL> name = global_properties [ SinkProperties . Label ] <EOL> self . _cache = deque ( ) <EOL> self . _waiters = deque ( ) <EOL> self . _min_size = sink_properties . min_watermark <EOL> self . _max_size = sink_properties . max_watermark <EOL> self . _max_queue_size = sink_properties . max_queue_len <EOL> self . _current_size = <NUM_LIT:0> <EOL> self . _state = ChannelState . Idle <EOL> socket_name = '<STR_LIT>' % ( endpoint . host , endpoint . port ) <EOL> self . endpoint = socket_name <EOL> self . _varz = self . Varz ( Source ( service = name , <EOL> endpoint = socket_name ) ) <EOL> self . _log = self . ROOT_LOG . getChild ( '<STR_LIT>' % ( name , socket_name ) ) <EOL> super ( WatermarkPoolSink , self ) . __init__ ( next_provider , global_properties ) <EOL> def __PropagateShutdown ( self , value ) : <EOL> self . on_faulted . Set ( value ) <EOL> def _DiscardSink ( self , sink ) : <EOL> """<STR_LIT>""" <EOL> sink . on_faulted . Unsubscribe ( self . __PropagateShutdown ) <EOL> sink . Close ( ) <EOL> def _Dequeue ( self ) : <EOL> """<STR_LIT>""" <EOL> while any ( self . _cache ) : <EOL> item = self . _cache . popleft ( ) <EOL> if item . state <= ChannelState . Open : <EOL> return item <EOL> else : <EOL> self . _DiscardSink ( item ) <EOL> return None <EOL> def _Get ( self ) : <EOL> cached = self . _Dequeue ( ) <EOL> if cached : <EOL> return cached <EOL> elif self . _current_size < self . _max_size : <EOL> self . _current_size += <NUM_LIT:1> <EOL> self . _varz . size ( self . _current_size ) <EOL> sink = self . _sink_provider . CreateSink ( self . _properties ) <EOL> sink . Open ( ) . wait ( ) <EOL> sink . on_faulted . Subscribe ( self . __PropagateShutdown ) <EOL> return sink <EOL> else : <EOL> if len ( self . _waiters ) + <NUM_LIT:1> > self . _max_queue_size : <EOL> return FailingMessageSink ( MaxWaitersError ( ) ) <EOL> else : <EOL> self . _varz . queue_size ( len ( self . _waiters ) + <NUM_LIT:1> ) <EOL> return QueuingMessageSink ( self . _waiters ) <EOL> def _Release ( self , sink ) : <EOL> if ( isinstance ( sink , QueuingMessageSink ) or <EOL> isinstance ( sink , FailingMessageSink ) ) : <EOL> self . _varz . queue_size ( len ( self . _waiters ) ) <EOL> return <EOL> do_close = False <EOL> if self . state == ChannelState . Closed : <EOL> self . _current_size -= <NUM_LIT:1> <EOL> elif sink . state == ChannelState . Closed : <EOL> self . _current_size -= <NUM_LIT:1> <EOL> self . Close ( ) <EOL> elif any ( self . _waiters ) : <EOL> gevent . spawn ( self . _ProcessQueue , sink ) <EOL> elif self . _current_size <= self . _min_size : <EOL> self . _cache . append ( sink ) <EOL> else : <EOL> self . _current_size -= <NUM_LIT:1> <EOL> do_close = True <EOL> self . _varz . size ( self . _current_size ) <EOL> if do_close : <EOL> self . _DiscardSink ( sink ) <EOL> def _ProcessQueue ( self , sink ) : <EOL> """<STR_LIT>""" <EOL> sink_stack , msg , stream , headers = self . _waiters . popleft ( ) <EOL> self . _varz . queue_size ( len ( self . _waiters ) ) <EOL> orig_sink , ctx = sink_stack . Pop ( ) <EOL> sink_stack . Push ( orig_sink , sink ) <EOL> sink . AsyncProcessRequest ( sink_stack , msg , stream , headers ) <EOL> def Open ( self ) : <EOL> ar = AsyncResult ( ) <EOL> ar . SafeLink ( self . _OpenImpl ) <EOL> return ar <EOL> def _OpenImpl ( self ) : <EOL> sink = self . _Get ( ) <EOL> self . _Release ( sink ) <EOL> self . _state = ChannelState . Open <EOL> def _FlushCache ( self ) : <EOL> [ self . _DiscardSink ( sink ) for sink in self . _cache ] <EOL> def Close ( self ) : <EOL> self . _state = ChannelState . Closed <EOL> self . _FlushCache ( ) <EOL> fail_sink = FailingMessageSink ( ServiceClosedError ) <EOL> [ fail_sink . AsyncProcessRequest ( sink_stack , msg , stream , headers ) <EOL> for sink_stack , msg , stream , headers in self . _waiters ] <EOL> @ property <EOL> def state ( self ) : <EOL> return self . _state <EOL> WatermarkPoolSink . Builder = SinkProvider ( <EOL> WatermarkPoolSink , <EOL> SinkRole . Pool , <EOL> min_watermark = <NUM_LIT:1> , <EOL> max_watermark = Int . MaxValue , <EOL> max_queue_len = Int . MaxValue <EOL> ) </s>
