1#!/usr/bin/env python
2# Copyright 2011 Google Inc. All Rights Reserved.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#      http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16"""Control "replay.py --server_mode" (e.g. switch from record to replay)."""
17
18import sys
19import time
20
21class ServerManager(object):
22  """Run servers until is removed or an exception is raised.
23
24  Servers start in the order they are appended and stop in the
25  opposite order. Servers are started by calling the initializer
26  passed to ServerManager.Append() and by calling __enter__(). Once an
27  server's initializer is called successfully, the __exit__() function
28  is guaranteed to be called when ServerManager.Run() completes.
29  """
30
31  def __init__(self, is_record_mode):
32    """Initialize a server manager."""
33    self.initializers = []
34    self.record_callbacks = []
35    self.replay_callbacks = []
36    self.traffic_shapers = []
37    self.is_record_mode = is_record_mode
38    self.should_exit = False
39
40  def Append(self, initializer, *init_args, **init_kwargs):
41    """Append a server to the end of the list to run.
42
43    Servers start in the order they are appended and stop in the
44    opposite order.
45
46    Args:
47      initializer: a function that returns a server instance.
48          A server needs to implement the with-statement interface.
49      init_args: positional arguments for the initializer.
50      init_args: keyword arguments for the initializer.
51    """
52    self.initializers.append((initializer, init_args, init_kwargs))
53
54  def AppendTrafficShaper(self, initializer, *init_args, **init_kwargs):
55    """Append a traffic shaper to the end of the list to run.
56
57    Args:
58      initializer: a function that returns a server instance.
59          A server needs to implement the with-statement interface.
60      init_args: positional arguments for the initializer.
61      init_args: keyword arguments for the initializer.
62    """
63    self.traffic_shapers.append((initializer, init_args, init_kwargs))
64
65  def AppendRecordCallback(self, func):
66    """Append a function to the list to call when switching to record mode.
67
68    Args:
69      func: a function that takes no arguments and returns no value.
70    """
71    self.record_callbacks.append(func)
72
73  def AppendReplayCallback(self, func):
74    """Append a function to the list to call when switching to replay mode.
75
76    Args:
77      func: a function that takes no arguments and returns no value.
78    """
79    self.replay_callbacks.append(func)
80
81  def IsRecordMode(self):
82    """Call all the functions that have been registered to enter replay mode."""
83    return self.is_record_mode
84
85  def SetRecordMode(self):
86    """Call all the functions that have been registered to enter record mode."""
87    self.is_record_mode = True
88    for record_func in self.record_callbacks:
89      record_func()
90
91  def SetReplayMode(self):
92    """Call all the functions that have been registered to enter replay mode."""
93    self.is_record_mode = False
94    for replay_func in self.replay_callbacks:
95      replay_func()
96
97  def Run(self):
98    """Create the servers and loop.
99
100    The loop quits if a server raises an exception.
101
102    Raises:
103      any exception raised by the servers
104    """
105    server_exits = []
106    server_ports = []
107    exception_info = (None, None, None)
108    try:
109      for initializer, init_args, init_kwargs in self.initializers:
110        server = initializer(*init_args, **init_kwargs)
111        if server:
112          server_exits.insert(0, server.__exit__)
113          server.__enter__()
114          if hasattr(server, 'server_port'):
115            server_ports.append(server.server_port)
116      for initializer, init_args, init_kwargs in self.traffic_shapers:
117        init_kwargs['ports'] = server_ports
118        shaper = initializer(*init_args, **init_kwargs)
119        if server:
120          server_exits.insert(0, shaper.__exit__)
121          shaper.__enter__()
122      while True:
123        time.sleep(1)
124        if self.should_exit:
125          break
126    except Exception:
127      exception_info = sys.exc_info()
128    finally:
129      for server_exit in server_exits:
130        try:
131          if server_exit(*exception_info):
132            exception_info = (None, None, None)
133        except Exception:
134          exception_info = sys.exc_info()
135      if exception_info != (None, None, None):
136        # pylint: disable=raising-bad-type
137        raise exception_info[0], exception_info[1], exception_info[2]
138