def test_evict_all(self): cache = StateCache(5) cache.put("key", "cache_token", "value") cache.put("key2", "cache_token", "value2") self.assertEqual(len(cache), 2) cache.evict_all() self.assertEqual(len(cache), 0) self.assertEqual(cache.get("key", "cache_token"), None) self.assertEqual(cache.get("key2", "cache_token"), None)
class GrpcStateHandlerFactory(StateHandlerFactory): """A factory for ``GrpcStateHandler``. Caches the created channels by ``state descriptor url``. """ def __init__(self, state_cache_size, credentials=None): self._state_handler_cache = {} self._lock = threading.Lock() self._throwing_state_handler = ThrowingStateHandler() self._credentials = credentials self._state_cache = StateCache(state_cache_size) def create_state_handler(self, api_service_descriptor): if not api_service_descriptor: return self._throwing_state_handler url = api_service_descriptor.url if url not in self._state_handler_cache: with self._lock: if url not in self._state_handler_cache: # Options to have no limits (-1) on the size of the messages # received or sent over the data plane. The actual buffer size is # controlled in a layer above. options = [('grpc.max_receive_message_length', -1), ('grpc.max_send_message_length', -1)] if self._credentials is None: logging.info('Creating insecure state channel for %s.', url) grpc_channel = GRPCChannelFactory.insecure_channel( url, options=options) else: logging.info('Creating secure state channel for %s.', url) grpc_channel = GRPCChannelFactory.secure_channel( url, self._credentials, options=options) logging.info('State channel established.') # Add workerId to the grpc channel grpc_channel = grpc.intercept_channel( grpc_channel, WorkerIdInterceptor()) self._state_handler_cache[ url] = CachingMaterializingStateHandler( self._state_cache, GrpcStateHandler( beam_fn_api_pb2_grpc.BeamFnStateStub( grpc_channel))) return self._state_handler_cache[url] def close(self): logging.info('Closing all cached gRPC state handlers.') for _, state_handler in self._state_handler_cache.items(): state_handler.done() self._state_handler_cache.clear() self._state_cache.evict_all()