1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
|
#!/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]
|
