Home | History | Annotate | Download | only in webpagereplay 
      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 
     18 import sys
     19 import time
     20 
     21 class 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