Permalink
Switch branches/tags
Find file
Fetching contributors…
Cannot retrieve contributors at this time
132 lines (109 sloc) 4.34 KB
#!/usr/bin/env python
# Copyright 2011 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Control "replay.py --server_mode" (e.g. switch from record to replay)."""
import sys
import time
class ServerManager(object):
"""Run servers until is removed or an exception is raised.
Servers start in the order they are appended and stop in the
opposite order. Servers are started by calling the initializer
passed to ServerManager.Append() and by calling __enter__(). Once an
server's initializer is called successfully, the __exit__() function
is guaranteed to be called when ServerManager.Run() completes.
"""
def __init__(self, is_record_mode):
"""Initialize a server manager."""
self.initializers = []
self.record_callbacks = []
self.replay_callbacks = []
self.is_record_mode = is_record_mode
def Append(self, initializer, *init_args, **init_kwargs):
"""Append a server to the end of the list to run.
Servers start in the order they are appended and stop in the
opposite order.
Args:
initializer: a function that returns a server instance.
A server needs to implement the with-statement interface.
init_args: positional arguments for the initializer.
init_args: keyword arguments for the initializer.
"""
self.initializers.append((initializer, init_args, init_kwargs))
def AppendStartStopFunctions(self, start_spec, stop_spec):
"""Append functions to call before and after the main run-loop.
If the enter function succeeds, then the exit function will be
called when shutting down.
Args:
start_spec: (start_func, start_args_1, start_arg_2, ...)
# The arguments are optional.
stop_spec: (stop_func, stop_args_1, stop_arg_2, ...)
# The arguments are optional.
"""
class Context(object):
def __enter__(self):
start_spec[0](*start_spec[1:])
def __exit__(self, type, value, traceback):
stop_spec[0](*stop_spec[1:])
self.Append(Context)
def AppendRecordCallback(self, func):
"""Append a function to the list to call when switching to record mode.
Args:
func: a function that takes no arguments and returns no value.
"""
self.record_callbacks.append(func)
def AppendReplayCallback(self, func):
"""Append a function to the list to call when switching to replay mode.
Args:
func: a function that takes no arguments and returns no value.
"""
self.replay_callbacks.append(func)
def IsRecordMode(self):
"""Call all the functions that have been registered to enter replay mode."""
return self.is_record_mode
def SetRecordMode(self):
"""Call all the functions that have been registered to enter record mode."""
self.is_record_mode = True
for record_func in self.record_callbacks:
record_func()
def SetReplayMode(self):
"""Call all the functions that have been registered to enter replay mode."""
self.is_record_mode = False
for replay_func in self.replay_callbacks:
replay_func()
def Run(self):
"""Create the servers and loop.
The loop quits if a server raises an exception.
Raises:
any exception raised by the servers
"""
server_exits = []
exception_info = (None, None, None)
try:
for initializer, init_args, init_kwargs in self.initializers:
server = initializer(*init_args, **init_kwargs)
server_exits.insert(0, server.__exit__)
server.__enter__()
while True:
time.sleep(1)
except:
exception_info = sys.exc_info()
finally:
for server_exit in server_exits:
try:
if server_exit(*exception_info):
exception_info = (None, None, None)
except:
exception_info = sys.exc_info()
if exception_info != (None, None, None):
raise exception_info[0], exception_info[1], exception_info[2]