servermanager.py

#!/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.traffic_shapers = []
    self.is_record_mode = is_record_mode
    self.should_exit = False

  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 AppendTrafficShaper(self, initializer, *init_args, **init_kwargs):
    """Append a traffic shaper to the end of the list to run.

    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.traffic_shapers.append((initializer, init_args, init_kwargs))

  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 = []
    server_ports = []
    exception_info = (None, None, None)
    try:
      for initializer, init_args, init_kwargs in self.initializers:
        server = initializer(*init_args, **init_kwargs)
        if server:
          server_exits.insert(0, server.__exit__)
          server.__enter__()
          if hasattr(server, 'server_port'):
            server_ports.append(server.server_port)
      for initializer, init_args, init_kwargs in self.traffic_shapers:
        init_kwargs['ports'] = server_ports
        shaper = initializer(*init_args, **init_kwargs)
        if server:
          server_exits.insert(0, shaper.__exit__)
          shaper.__enter__()
      while True:
        time.sleep(1)
        if self.should_exit:
          break
    except Exception:
      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:
          exception_info = sys.exc_info()
      if exception_info != (None, None, None):
        # pylint: disable=raising-bad-type
        raise exception_info[0], exception_info[1], exception_info[2]