mirror of https://github.com/grpc/grpc.git
The C based gRPC (C++, Python, Ruby, Objective-C, PHP, C#)
https://grpc.io/
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
158 lines
5.8 KiB
158 lines
5.8 KiB
8 years ago
|
import struct
|
||
|
import messages_pb2
|
||
|
import functools
|
||
|
import argparse
|
||
|
import logging
|
||
|
import time
|
||
|
|
||
|
from twisted.internet.defer import Deferred, inlineCallbacks
|
||
|
from twisted.internet.protocol import Protocol, Factory
|
||
|
from twisted.internet import endpoints, reactor, error, defer
|
||
|
from h2.connection import H2Connection
|
||
|
from h2.events import RequestReceived, DataReceived, WindowUpdated, RemoteSettingsChanged, PingAcknowledged
|
||
|
from threading import Lock
|
||
|
|
||
|
READ_CHUNK_SIZE = 16384
|
||
|
GRPC_HEADER_SIZE = 5
|
||
|
|
||
|
class H2ProtocolBaseServer(Protocol):
|
||
|
def __init__(self):
|
||
|
self._conn = H2Connection(client_side=False)
|
||
|
self._recv_buffer = ''
|
||
|
self._handlers = {}
|
||
|
self._handlers['DataReceived'] = self.on_data_received_default
|
||
|
self._handlers['WindowUpdated'] = self.on_window_update_default
|
||
|
self._handlers['RequestReceived'] = self.on_request_received_default
|
||
|
self._handlers['SendDone'] = self.on_send_done_default
|
||
|
self._handlers['ConnectionLost'] = self.on_connection_lost
|
||
|
self._handlers['PingAcknowledged'] = self.on_ping_acknowledged_default
|
||
|
self._stream_status = {}
|
||
|
self._outstanding_pings = 0
|
||
|
|
||
|
def set_handlers(self, handlers):
|
||
|
self._handlers = handlers
|
||
|
|
||
|
def connectionMade(self):
|
||
|
logging.info('Connection Made')
|
||
|
self._conn.initiate_connection()
|
||
|
self.transport.setTcpNoDelay(True)
|
||
|
self.transport.write(self._conn.data_to_send())
|
||
|
|
||
|
def connectionLost(self, reason):
|
||
|
self._handlers['ConnectionLost'](reason)
|
||
|
|
||
|
def on_connection_lost(self, reason):
|
||
|
logging.info('Disconnected %s'%reason)
|
||
|
reactor.callFromThread(reactor.stop)
|
||
|
|
||
|
def dataReceived(self, data):
|
||
|
events = self._conn.receive_data(data)
|
||
|
if self._conn.data_to_send:
|
||
|
self.transport.write(self._conn.data_to_send())
|
||
|
for event in events:
|
||
|
if isinstance(event, RequestReceived) and self._handlers.has_key('RequestReceived'):
|
||
|
logging.info('RequestReceived Event')
|
||
|
self._handlers['RequestReceived'](event)
|
||
|
elif isinstance(event, DataReceived) and self._handlers.has_key('DataReceived'):
|
||
|
logging.info('DataReceived Event')
|
||
|
self._handlers['DataReceived'](event)
|
||
|
elif isinstance(event, WindowUpdated) and self._handlers.has_key('WindowUpdated'):
|
||
|
logging.info('WindowUpdated Event')
|
||
|
self._handlers['WindowUpdated'](event)
|
||
|
elif isinstance(event, PingAcknowledged) and self._handlers.has_key('PingAcknowledged'):
|
||
|
logging.info('PingAcknowledged Event')
|
||
|
self._handlers['PingAcknowledged'](event)
|
||
|
self.transport.write(self._conn.data_to_send())
|
||
|
|
||
|
def on_ping_acknowledged_default(self, event):
|
||
|
self._outstanding_pings -= 1
|
||
|
|
||
|
def on_data_received_default(self, event):
|
||
|
self._conn.acknowledge_received_data(len(event.data), event.stream_id)
|
||
|
self._recv_buffer += event.data
|
||
|
|
||
|
def on_request_received_default(self, event):
|
||
|
self._recv_buffer = ''
|
||
|
self._stream_id = event.stream_id
|
||
|
self._stream_status[event.stream_id] = True
|
||
|
self._conn.send_headers(
|
||
|
stream_id=event.stream_id,
|
||
|
headers=[
|
||
|
(':status', '200'),
|
||
|
('content-type', 'application/grpc'),
|
||
|
('grpc-encoding', 'identity'),
|
||
|
('grpc-accept-encoding', 'identity,deflate,gzip'),
|
||
|
],
|
||
|
)
|
||
|
self.transport.write(self._conn.data_to_send())
|
||
|
|
||
|
def on_window_update_default(self, event):
|
||
|
pass
|
||
|
|
||
|
def send_reset_stream(self):
|
||
|
self._conn.reset_stream(self._stream_id)
|
||
|
self.transport.write(self._conn.data_to_send())
|
||
|
|
||
|
def setup_send(self, data_to_send):
|
||
|
self._send_remaining = len(data_to_send)
|
||
|
self._send_offset = 0
|
||
|
self._data_to_send = data_to_send
|
||
|
self.default_send()
|
||
|
|
||
|
def default_send(self):
|
||
|
while self._send_remaining > 0:
|
||
|
lfcw = self._conn.local_flow_control_window(self._stream_id)
|
||
|
if lfcw == 0:
|
||
|
break
|
||
|
chunk_size = min(lfcw, READ_CHUNK_SIZE)
|
||
|
bytes_to_send = min(chunk_size, self._send_remaining)
|
||
|
logging.info('flow_control_window = %d. sending [%d:%d] stream_id %d'%
|
||
|
(lfcw, self._send_offset, self._send_offset + bytes_to_send,
|
||
|
self._stream_id))
|
||
|
data = self._data_to_send[self._send_offset : self._send_offset + bytes_to_send]
|
||
|
self._conn.send_data(self._stream_id, data, False)
|
||
|
self._send_remaining -= bytes_to_send
|
||
|
self._send_offset += bytes_to_send
|
||
|
if self._send_remaining == 0:
|
||
|
self._handlers['SendDone']()
|
||
|
|
||
|
def default_ping(self):
|
||
|
self._outstanding_pings += 1
|
||
|
self._conn.ping(b'\x00'*8)
|
||
|
self.transport.write(self._conn.data_to_send())
|
||
|
|
||
|
def on_send_done_default(self):
|
||
|
if self._stream_status[self._stream_id]:
|
||
|
self._stream_status[self._stream_id] = False
|
||
|
self.default_send_trailer()
|
||
|
|
||
|
def default_send_trailer(self):
|
||
|
logging.info('Sending trailer for stream id %d'%self._stream_id)
|
||
|
self._conn.send_headers(self._stream_id,
|
||
|
headers=[ ('grpc-status', '0') ],
|
||
|
end_stream=True
|
||
|
)
|
||
|
self.transport.write(self._conn.data_to_send())
|
||
|
|
||
|
@staticmethod
|
||
|
def default_response_data(response_size):
|
||
|
sresp = messages_pb2.SimpleResponse()
|
||
|
sresp.payload.body = b'\x00'*response_size
|
||
|
serialized_resp_proto = sresp.SerializeToString()
|
||
|
response_data = b'\x00' + struct.pack('i', len(serialized_resp_proto))[::-1] + serialized_resp_proto
|
||
|
return response_data
|
||
|
|
||
|
@staticmethod
|
||
|
def parse_received_data(recv_buffer):
|
||
|
""" returns a grpc framed string of bytes containing response proto of the size
|
||
|
asked in request """
|
||
|
grpc_msg_size = struct.unpack('i',recv_buffer[1:5][::-1])[0]
|
||
|
if len(recv_buffer) != GRPC_HEADER_SIZE + grpc_msg_size:
|
||
|
logging.error('not enough data to decode req proto. size = %d, needed %s'%(len(recv_buffer), 5+grpc_msg_size))
|
||
|
return None
|
||
|
req_proto_str = recv_buffer[5:5+grpc_msg_size]
|
||
|
sr = messages_pb2.SimpleRequest()
|
||
|
sr.ParseFromString(req_proto_str)
|
||
|
logging.info('Parsed request: response_size=%s'%sr.response_size)
|
||
|
return sr
|