blob: 74933cd009a9ae64819eeac7d561b342d46d0cb7 [file] [log] [blame]
Zsolt Harasztia9a12dc2016-09-27 13:48:35 -07001#
2# Copyright 2016 the original author or authors.
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
17"""
Zsolt Harasztie7b60762016-10-05 17:49:27 -070018gRPC client meant to connect to a gRPC server endpoint, and query the
Zsolt Haraszticba96de2016-11-06 14:04:55 -080019end-point's schema by calling SchemaService.Schema(Empty) and all of its
Zsolt Harasztia9a12dc2016-09-27 13:48:35 -070020semantics are derived from the recovered schema.
21"""
Zsolt Harasztie7b60762016-10-05 17:49:27 -070022
Zsolt Harasztia9a12dc2016-09-27 13:48:35 -070023import os
Zsolt Harasztia9a12dc2016-09-27 13:48:35 -070024import sys
Zsolt Harasztia9a12dc2016-09-27 13:48:35 -070025from random import randint
Zsolt Harasztia9a12dc2016-09-27 13:48:35 -070026from zlib import decompress
27
Zsolt Haraszti3d55ffc2016-10-03 22:26:41 -070028import grpc
29from consul import Consul
Zsolt Harasztie7b60762016-10-05 17:49:27 -070030from grpc._channel import _Rendezvous
Zsolt Haraszti3d55ffc2016-10-03 22:26:41 -070031from structlog import get_logger
Zsolt Harasztie7b60762016-10-05 17:49:27 -070032from twisted.internet import reactor
Zsolt Haraszti2aac6232016-11-23 11:18:23 -080033from twisted.internet.defer import inlineCallbacks, returnValue
Zsolt Harasztie7b60762016-10-05 17:49:27 -070034from werkzeug.exceptions import ServiceUnavailable
Zsolt Haraszti3d55ffc2016-10-03 22:26:41 -070035
Zsolt Harasztiabae5912016-10-16 19:30:34 -070036from common.utils.asleep import asleep
Zsolt Haraszti21980762016-11-08 10:57:19 -080037from chameleon.protos import third_party
Zsolt Haraszticba96de2016-11-06 14:04:55 -080038from chameleon.protos.schema_pb2 import SchemaServiceStub
39from google.protobuf.empty_pb2 import Empty
40
Zsolt Harasztia9a12dc2016-09-27 13:48:35 -070041
42log = get_logger()
43
44
45class GrpcClient(object):
Zsolt Harasztie7b60762016-10-05 17:49:27 -070046 """
47 Connect to a gRPC server, fetch its schema, and process the downloaded
48 schema files to drive the customization of the north-bound interface(s)
49 of Chameleon.
50 """
51 RETRY_BACKOFF = [0.05, 0.1, 0.2, 0.5, 1, 2, 5]
Zsolt Harasztia9a12dc2016-09-27 13:48:35 -070052
Zsolt Harasztidca6fa12016-11-03 16:56:17 -070053 def __init__(self, consul_endpoint, work_dir, endpoint='localhost:50055',
54 reconnect_callback=None):
Zsolt Harasztia9a12dc2016-09-27 13:48:35 -070055 self.consul_endpoint = consul_endpoint
56 self.endpoint = endpoint
Zsolt Haraszti3d55ffc2016-10-03 22:26:41 -070057 self.work_dir = work_dir
Zsolt Harasztidca6fa12016-11-03 16:56:17 -070058 self.reconnect_callback = reconnect_callback
59
Zsolt Haraszti3d55ffc2016-10-03 22:26:41 -070060 self.plugin_dir = os.path.abspath(os.path.join(
61 os.path.dirname(__file__), '../protoc_plugins'))
Zsolt Harasztia9a12dc2016-09-27 13:48:35 -070062
63 self.channel = None
64 self.schema = None
Zsolt Harasztie7b60762016-10-05 17:49:27 -070065 self.retries = 0
Zsolt Harasztia9a12dc2016-09-27 13:48:35 -070066 self.shutting_down = False
Zsolt Haraszti3cf36342016-10-05 20:40:19 -070067 self.connected = False
Zsolt Harasztia9a12dc2016-09-27 13:48:35 -070068
Zsolt Harasztidca6fa12016-11-03 16:56:17 -070069 def start(self):
70 log.debug('starting')
Zsolt Haraszti3cf36342016-10-05 20:40:19 -070071 if not self.connected:
72 reactor.callLater(0, self.connect)
Zsolt Harasztidca6fa12016-11-03 16:56:17 -070073 log.info('started')
Zsolt Harasztia9a12dc2016-09-27 13:48:35 -070074 return self
75
Zsolt Harasztidca6fa12016-11-03 16:56:17 -070076 def stop(self):
77 log.debug('stopping')
Zsolt Harasztia9a12dc2016-09-27 13:48:35 -070078 if self.shutting_down:
79 return
80 self.shutting_down = True
Zsolt Harasztidca6fa12016-11-03 16:56:17 -070081 log.info('stopped')
82
83 def set_reconnect_callback(self, reconnect_callback):
84 self.reconnect_callback = reconnect_callback
85 return self
Zsolt Harasztia9a12dc2016-09-27 13:48:35 -070086
Zsolt Harasztie7b60762016-10-05 17:49:27 -070087 @inlineCallbacks
Zsolt Harasztia9a12dc2016-09-27 13:48:35 -070088 def connect(self):
Zsolt Harasztie7b60762016-10-05 17:49:27 -070089 """
90 (Re-)Connect to end-point
91 """
92
Zsolt Haraszti3cf36342016-10-05 20:40:19 -070093 if self.shutting_down or self.connected:
Zsolt Harasztia9a12dc2016-09-27 13:48:35 -070094 return
95
96 try:
97 if self.endpoint.startswith('@'):
Zsolt Harasztie7b60762016-10-05 17:49:27 -070098 _endpoint = self._get_endpoint_from_consul(self.endpoint[1:])
Zsolt Harasztia9a12dc2016-09-27 13:48:35 -070099 else:
100 _endpoint = self.endpoint
101
102 log.info('connecting', endpoint=_endpoint)
103 self.channel = grpc.insecure_channel(_endpoint)
104
Zsolt Haraszti21980762016-11-08 10:57:19 -0800105 swagger_from = self._retrieve_schema()
106 self._compile_proto_files(swagger_from)
Zsolt Harasztie7b60762016-10-05 17:49:27 -0700107 self._clear_backoff()
108
Zsolt Haraszti3cf36342016-10-05 20:40:19 -0700109 self.connected = True
Zsolt Harasztidca6fa12016-11-03 16:56:17 -0700110 if self.reconnect_callback is not None:
111 reactor.callLater(0, self.reconnect_callback)
Zsolt Harasztie7b60762016-10-05 17:49:27 -0700112
113 return
114
115 except _Rendezvous, e:
116 if e.code() == grpc.StatusCode.UNAVAILABLE:
117 log.info('grpc-endpoint-not-available')
118 else:
119 log.exception(e)
120 yield self._backoff('not-available')
Zsolt Harasztia9a12dc2016-09-27 13:48:35 -0700121
122 except Exception, e:
Zsolt Harasztie7b60762016-10-05 17:49:27 -0700123 if not self.shutting_down:
124 log.exception('cannot-connect', endpoint=_endpoint)
125 yield self._backoff('unknown-error')
Zsolt Harasztia9a12dc2016-09-27 13:48:35 -0700126
Zsolt Harasztie7b60762016-10-05 17:49:27 -0700127 reactor.callLater(0, self.connect)
128
129 def _backoff(self, msg):
130 wait_time = self.RETRY_BACKOFF[min(self.retries,
131 len(self.RETRY_BACKOFF) - 1)]
132 self.retries += 1
133 log.error(msg, retry_in=wait_time)
134 return asleep(wait_time)
135
136 def _clear_backoff(self):
137 if self.retries:
138 log.info('reconnected', after_retries=self.retries)
139 self.retries = 0
140
141 def _get_endpoint_from_consul(self, service_name):
142 """
143 Look up an appropriate grpc endpoint (host, port) from
144 consul, under the service name specified by service-name
Zsolt Harasztia9a12dc2016-09-27 13:48:35 -0700145 """
146 host = self.consul_endpoint.split(':')[0].strip()
147 port = int(self.consul_endpoint.split(':')[1].strip())
148
149 consul = Consul(host=host, port=port)
150 _, services = consul.catalog.service(service_name)
151
152 if len(services) == 0:
153 raise Exception('Cannot find service %s in consul' % service_name)
154
155 # pick a random entry
156 # TODO should we prefer local IP addresses? Probably.
157
158 service = services[randint(0, len(services) - 1)]
159 endpoint = '{}:{}'.format(service['ServiceAddress'],
160 service['ServicePort'])
161 return endpoint
162
163 def _retrieve_schema(self):
Zsolt Harasztie7b60762016-10-05 17:49:27 -0700164 """
165 Retrieve schema from gRPC end-point, and save all *.proto files in
166 the work directory.
167 """
Zsolt Harasztia9a12dc2016-09-27 13:48:35 -0700168 assert isinstance(self.channel, grpc.Channel)
169 stub = SchemaServiceStub(self.channel)
Zsolt Harasztie7b60762016-10-05 17:49:27 -0700170 # try:
Zsolt Haraszticba96de2016-11-06 14:04:55 -0800171 schemas = stub.GetSchema(Empty())
Zsolt Harasztie7b60762016-10-05 17:49:27 -0700172 # except _Rendezvous, e:
173 # if e.code == grpc.StatusCode.UNAVAILABLE:
174 #
175 # else:
176 # raise e
Zsolt Harasztia9a12dc2016-09-27 13:48:35 -0700177
178 os.system('mkdir -p %s' % self.work_dir)
179 os.system('rm -fr /tmp/%s/*' %
180 self.work_dir.replace('/tmp/', '')) # safer
181
Zsolt Harasztidca6fa12016-11-03 16:56:17 -0700182 for proto_file in schemas.protos:
183 proto_fname = proto_file.file_name
184 proto_content = proto_file.proto
185 log.debug('saving-proto', fname=proto_fname, dir=self.work_dir,
186 length=len(proto_content))
187 with open(os.path.join(self.work_dir, proto_fname), 'w') as f:
188 f.write(proto_content)
Zsolt Harasztia9a12dc2016-09-27 13:48:35 -0700189
Zsolt Harasztidca6fa12016-11-03 16:56:17 -0700190 desc_content = decompress(proto_file.descriptor)
191 desc_fname = proto_fname.replace('.proto', '.desc')
192 log.debug('saving-descriptor', fname=desc_fname, dir=self.work_dir,
193 length=len(desc_content))
194 with open(os.path.join(self.work_dir, desc_fname), 'wb') as f:
195 f.write(desc_content)
Zsolt Haraszti21980762016-11-08 10:57:19 -0800196 return schemas.swagger_from
Zsolt Harasztia9a12dc2016-09-27 13:48:35 -0700197
Zsolt Haraszti21980762016-11-08 10:57:19 -0800198 def _compile_proto_files(self, swagger_from):
Zsolt Harasztie7b60762016-10-05 17:49:27 -0700199 """
200 For each *.proto file in the work directory, compile the proto
201 file into the respective *_pb2.py file as well as generate the
202 web server gateway python file *_gw.py.
203 :return: None
204 """
Zsolt Harasztia9a12dc2016-09-27 13:48:35 -0700205 google_api_dir = os.path.abspath(os.path.join(
Zsolt Haraszti46c72002016-10-10 09:55:30 -0700206 os.path.dirname(__file__), '../protos/third_party'
207 ))
208
209 chameleon_base_dir = os.path.abspath(os.path.join(
210 os.path.dirname(__file__), '../..'
Zsolt Harasztia9a12dc2016-09-27 13:48:35 -0700211 ))
212
213 for fname in [f for f in os.listdir(self.work_dir)
214 if f.endswith('.proto')]:
Zsolt Harasztia9a12dc2016-09-27 13:48:35 -0700215
Zsolt Haraszti21980762016-11-08 10:57:19 -0800216 need_swagger = fname == swagger_from
217 log.debug('compiling', file=fname, need_swagger=need_swagger)
Zsolt Harasztia9a12dc2016-09-27 13:48:35 -0700218 cmd = (
219 'cd %s && '
Zsolt Haraszti05b837a2016-10-05 00:18:57 -0700220 'env PATH=%s PYTHONPATH=%s '
Zsolt Harasztia9a12dc2016-09-27 13:48:35 -0700221 'python -m grpc.tools.protoc '
222 '-I. '
223 '-I%s '
224 '--python_out=. '
225 '--grpc_python_out=. '
Zsolt Haraszti3d55ffc2016-10-03 22:26:41 -0700226 '--plugin=protoc-gen-gw=%s/gw_gen.py '
227 '--gw_out=. '
Zsolt Haraszti46c72002016-10-10 09:55:30 -0700228 '--plugin=protoc-gen-swagger=%s/swagger_gen.py '
Zsolt Haraszti21980762016-11-08 10:57:19 -0800229 '%s'
Zsolt Haraszti3d55ffc2016-10-03 22:26:41 -0700230 '%s' % (
231 self.work_dir,
232 ':'.join([os.environ['PATH'], self.plugin_dir]),
Zsolt Haraszti46c72002016-10-10 09:55:30 -0700233 ':'.join([google_api_dir, chameleon_base_dir]),
Zsolt Haraszti3d55ffc2016-10-03 22:26:41 -0700234 google_api_dir,
Zsolt Haraszti46c72002016-10-10 09:55:30 -0700235 self.plugin_dir,
Zsolt Haraszti3d55ffc2016-10-03 22:26:41 -0700236 self.plugin_dir,
Zsolt Haraszti21980762016-11-08 10:57:19 -0800237 '--swagger_out=. ' if need_swagger else '',
Zsolt Haraszti3d55ffc2016-10-03 22:26:41 -0700238 fname)
Zsolt Harasztia9a12dc2016-09-27 13:48:35 -0700239 )
Zsolt Harasztie7b60762016-10-05 17:49:27 -0700240 log.debug('executing', cmd=cmd, file=fname)
Zsolt Harasztia9a12dc2016-09-27 13:48:35 -0700241 os.system(cmd)
Zsolt Harasztie7b60762016-10-05 17:49:27 -0700242 log.info('compiled', file=fname)
Zsolt Harasztia9a12dc2016-09-27 13:48:35 -0700243
244 # test-load each _pb2 file to see all is right
245 if self.work_dir not in sys.path:
246 sys.path.insert(0, self.work_dir)
247
248 for fname in [f for f in os.listdir(self.work_dir)
249 if f.endswith('_pb2.py')]:
250 modname = fname[:-len('.py')]
251 log.debug('test-import', modname=modname)
252 _ = __import__(modname)
Zsolt Haraszti3d55ffc2016-10-03 22:26:41 -0700253
Zsolt Haraszti2aac6232016-11-23 11:18:23 -0800254 @inlineCallbacks
255 def invoke(self, stub, method_name, request, retry=1):
Zsolt Harasztie7b60762016-10-05 17:49:27 -0700256 """
257 Invoke a gRPC call to the remote server and return the response.
258 :param stub: Reference to the *_pb2 service stub
259 :param method_name: The method name inside the service stub
260 :param request: The request protobuf message
261 :return: The response protobuf message
262 """
263
Zsolt Haraszti3cf36342016-10-05 20:40:19 -0700264 if not self.connected:
Zsolt Harasztie7b60762016-10-05 17:49:27 -0700265 raise ServiceUnavailable()
266
267 try:
268 response = getattr(stub(self.channel), method_name)(request)
Zsolt Haraszti2aac6232016-11-23 11:18:23 -0800269 returnValue(response)
Zsolt Harasztie7b60762016-10-05 17:49:27 -0700270
271 except grpc._channel._Rendezvous, e:
272 if e.code() == grpc.StatusCode.UNAVAILABLE:
273 e = ServiceUnavailable()
Zsolt Haraszti2aac6232016-11-23 11:18:23 -0800274
275 if self.connected:
276 self.connected = False
277 yield self.connect()
278 if retry > 0:
279 response = yield self.invoke(stub, method_name,
280 request,
281 retry=retry - 1)
282 returnValue(response)
283
Zsolt Harasztie7b60762016-10-05 17:49:27 -0700284 else:
285 log.exception(e)
286
Zsolt Harasztie7b60762016-10-05 17:49:27 -0700287 raise e