def from_bytes(klass, data): """Returns thrift object from a string, using standard binary representation.""" obj = klass() b = TMemoryBuffer(data) p = TBinaryProtocol(b) obj.read(p) return obj
def make_broad_cast_node_info(node_id, name):
m = TMemoryBuffer()
p = TBinaryProtocol(m)
msg_name = SERVICE_BROADCAST_REQ
nonce = get_random_id()
head = msg_header(get_magic(), msg_name, nonce)
head.write(p)
node_info = node_service_info()
node_info.name = name
node_info.time_stamp = int(time.time())
service_list = []
service_list.append("ai_training")
node_info.service_list = service_list
kvs = {}
kvs["gpu"] = "1 * GeForce940MX"
kvs["state"] = "idle"
node_info.kvs = kvs
node_map = {}
node_map[node_id] = node_info
req = service_broadcast_req_body(node_map)
req.write(p)
p.writeMessageEnd()
m.flush()
return pack_head(m)
def test_with_headers(self):
client_memory_trans = TMemoryBuffer()
client_prot = THeaderProtocol(client_memory_trans)
client_header_trans = client_prot.trans
client_header_trans.set_header("Trace", "1234")
client_header_trans.set_header("Parent", "2345")
client_header_trans.set_header("Span", "3456")
client = BaseplateService.Client(client_prot)
try:
client.is_healthy()
except:
pass # we don't have a test response for the client
self.itrans._readBuffer = StringIO(client_memory_trans.getvalue())
self.processor.process(self.iprot, self.oprot, self.server_context)
self.assertEqual(self.observer.on_server_span_created.call_count, 1)
context, server_span = self.observer.on_server_span_created.call_args[0]
self.assertEqual(server_span.trace_id, 1234)
self.assertEqual(server_span.parent_id, 2345)
self.assertEqual(server_span.id, 3456)
self.assertTrue(self.server_observer.on_start.called)
self.assertTrue(self.server_observer.on_finish.called)
def test_primitive_serialization(self):
val = int(12)
buf = TMemoryBuffer()
proto = TCompactProtocol(buf)
proto.writeI32(val)
reader = TCompactProtocol(StringIO(buf.getvalue()))
self.assertEqual(reader.readI32(), val)
async def _on_message_callback(self, message):
"""The function to be executed when a message is received."""
if not message.reply:
logger.warn('no reply present, discarding message')
return
frame_size = struct.unpack('!I', message.data[:4])[0]
if frame_size > _NATS_MAX_MESSAGE_SIZE - 4:
logger.warning('frame size too large, dropping message')
return
# process frame, first four bytes are the frame size
iprot = self._protocol_factory.get_protocol(
TMemoryBuffer(message.data[4:]))
otrans = TMemoryOutputBuffer(_NATS_MAX_MESSAGE_SIZE)
oprot = self._protocol_factory.get_protocol(otrans)
try:
await self._processor.process(iprot, oprot)
except TApplicationException:
# Continue so the exception is sent to the client
pass
except Exception:
return
if len(otrans) == 4:
return
await self._nats_client.publish(message.reply, otrans.getvalue())
async def request(self, context: FContext, payload) -> TTransportBase:
"""
Write the current buffer payload over the network and return the
response.
"""
self._preflight_request_check(payload)
encoded = base64.b64encode(payload)
status, text = await self._make_request(context, encoded)
if status == 413:
raise TTransportException(
type=TTransportExceptionType.RESPONSE_TOO_LARGE,
message='response was too large for the transport')
if status >= 300:
raise TTransportException(
type=TTransportExceptionType.UNKNOWN,
message='request errored with code {0} and message {1}'.format(
status, str(text)))
decoded = base64.b64decode(text)
if len(decoded) < 4:
raise TTransportException(type=TTransportExceptionType.UNKNOWN,
message='invalid frame size')
if len(decoded) == 4:
if any(decoded):
raise TTransportException(type=TTransportExceptionType.UNKNOWN,
message='missing data')
# One-way method, drop response
return
return TMemoryBuffer(decoded[4:])
def request(self, context, payload):
self._preflight_request_check(payload)
op_id = str(context._get_op_id())
future = gen.Future()
with (yield self._futures_lock.acquire()):
if op_id in self._futures:
raise TTransportException(
type=TTransportExceptionType.UNKNOWN,
message="request already in flight for context")
self._futures[op_id] = future
try:
@gen.coroutine
def flush_and_wait():
yield self.flush(payload)
data = yield future
raise gen.Return(data)
data = yield gen.with_timeout(
timedelta(milliseconds=context.timeout), flush_and_wait())
raise gen.Return(TMemoryBuffer(data))
except gen.TimeoutError:
raise TTransportException(type=TTransportExceptionType.TIMED_OUT,
message='request timed out')
finally:
with (yield self._futures_lock.acquire()):
del self._futures[op_id]
def get_doc_hot_by_thread(docid_list):
docid_hot_list = []
r = redis.Redis(host='***', port=1680, password='******', charset='gbk')
redis_result_list = []
try:
redis_result_list = r.mget(docid_list)
if redis_result_list is not None and len(redis_result_list) == len(
docid_list):
for i in range(len(redis_result_list)):
if redis_result_list[i] is None or len(
redis_result_list[i]) == 0:
continue
docid = docid_list[i]
tMemory_o = TMemoryBuffer(redis_result_list[i])
tBinaryProtocol_o = TBinaryProtocol(tMemory_o)
hot_feature = HotFeature()
hot_feature.read(tBinaryProtocol_o)
pv = hot_feature.app_show_num if hot_feature.app_show_num is not None else 0
cl = hot_feature.app_read_num if hot_feature.app_read_num is not None else 0
readtime = hot_feature.app_read_duration_double if hot_feature.app_read_duration_double is not None else 0
coldstart_pv = hot_feature.coldstart_show_num if hot_feature.coldstart_show_num is not None else 0
coldstart_cl = hot_feature.coldstart_read_num if hot_feature.coldstart_read_num is not None else 0
docid_hot_list.append(
(docid, pv, cl, readtime, coldstart_pv, coldstart_cl))
except:
traceback.print_exc()
return docid_hot_list
def readStructFromJSON(jstr, struct_type):
stuff = struct_type()
trans = TMemoryBuffer(jstr)
proto = TSimpleJSONProtocol.TSimpleJSONProtocol(trans,
struct_type.thrift_spec)
stuff.read(proto)
return stuff
def make_get_peer_nodes_resp(dport):
m = TMemoryBuffer()
p = TBinaryProtocol(m)
msg_name = P2P_GET_PEER_NODES_RESP
nonce = get_random_id()
head = msg_header(get_magic(), msg_name, nonce)
head.write(p)
print("nonce:%s" % (nonce))
# self.peer_node_id = peer_node_id
# self.core_version = core_version
# self.protocol_version = protocol_version
# self.live_time_stamp = live_time_stamp
# self.addr = addr
node_list = []
node_info = peer_node_info()
node_info.addr = network_address("10.10.254.198", dport)
node_info.peer_node_id = gen_node_id()
node_info.core_version = core_version
node_info.protocol_version = pro_version
node_info.live_time_stamp = int(time.time())
node_list.append(node_info)
req = get_peer_nodes_resp_body(node_list)
req.write(p)
p.writeMessageEnd()
m.flush()
return pack_head(m)
def make_stop_training_req(in_task_id):
m = TMemoryBuffer()
p = TBinaryProtocol(m)
msg_name = STOP_TRAINING_REQ
nonce = get_random_id()
session_id = get_random_id()
head = msg_header(get_magic(), msg_name, nonce, session_id)
print("stop req.nonce:%s, task_id:%s" %(nonce, in_task_id))
req=stop_training_req_body(in_task_id)
message = in_task_id + nonce
sign_algo = "ecdsa"
origin = get_node_id()
exten_info = {}
exten_info["origin_id"] = origin
exten_info["sign_algo"] = sign_algo
exten_info["sign"] = dbc_sign(message)
head.exten_info = exten_info
head.write(p)
req.write(p)
p.writeMessageEnd()
m.flush()
return pack_head(m)
def write(self, buf):
"""Bounded write to buffer"""
if len(self) + len(buf) > self._limit > 0:
self._buffer = TMemoryBuffer()
raise TTransportException(
type=TTransportExceptionType.REQUEST_TOO_LARGE,
message="Buffer size reached {}".format(self._limit))
self._buffer.write(buf)
def to_LinkAttr(self, body):
link_info = LinkAttr()
try:
tMemory_o = TMemoryBuffer(body)
tBinaryProtocol_o = TBinaryProtocol(tMemory_o)
link_info.read(tBinaryProtocol_o)
except:
print traceback.format_exc()
def test_unexpected_exception_passed_to_server_span_finish(self):
client_memory_trans = TMemoryBuffer()
client_prot = THeaderProtocolFactory().getProtocol(client_memory_trans)
client = TestService.Client(client_prot)
try:
client.example_throws(crash=True)
except (TTransportException, EOFError):
pass # we don't have a test response for the client
itrans = TMemoryBuffer(client_memory_trans.getvalue())
iprot = THeaderProtocolFactory().getProtocol(itrans)
self.processor.process(iprot, self.oprot)
self.assertEqual(self.server_observer.on_start.call_count, 1)
self.assertEqual(self.server_observer.on_finish.call_count, 1)
_, captured_exc, _ = self.server_observer.on_finish.call_args[0][0]
self.assertIsInstance(captured_exc, UnexpectedException)
def test_tree():
t = branch(branch(
leaf("aa"),
leaf("ab"),
), branch(
leaf("ba"),
leaf("bb"),
))
transport_out = TMemoryBuffer()
protocol_out = TJSONProtocol(transport_out)
t.write(protocol_out)
json_out = transport_out.getvalue()
transport_in = TMemoryBuffer(json_out)
protocol_in = TJSONProtocol(transport_in)
u = Tree()
u.read(protocol_in)
assert t == u
def service_core_json(request):
try:
iprot = TJSONProtocol.TJSONProtocol(TMemoryBuffer(request.body))
oprot = TJSONProtocol.TJSONProtocol(TMemoryBuffer())
core_processor.process(iprot, oprot)
resp = HttpResponse(oprot.trans.getvalue(),
content_type="application/json")
try:
resp['Access-Control-Allow-Origin'] = request.META['HTTP_ORIGIN']
except:
pass
return resp
except Exception, e:
return HttpResponse(json.dumps({'error': str(e)}),
content_type="application/json")
def test_client_proxy_flow(self):
client_memory_trans = TMemoryBuffer()
client_prot = THeaderProtocolFactory().getProtocol(client_memory_trans)
class Pool(object):
@contextlib.contextmanager
def connection(self):
yield client_prot
client_factory = ThriftContextFactory(Pool(), TestService.Client)
span = mock.MagicMock()
child_span = span.make_child().__enter__()
child_span.trace_id = 1
child_span.parent_id = 1
child_span.id = 1
child_span.sampled = True
child_span.flags = None
edge_context = self.edge_context_factory.from_upstream(
SERIALIZED_EDGECONTEXT_WITH_VALID_AUTH)
edge_context.attach_context(child_span.context)
client = client_factory.make_object_for_context("test", span)
try:
client.example_simple()
except (TTransportException, EOFError):
pass # we don't have a test response for the client
itrans = TMemoryBuffer(client_memory_trans.getvalue())
iprot = THeaderProtocolFactory().getProtocol(itrans)
self.processor.process(iprot, self.oprot)
context, _ = self.observer.on_server_span_created.call_args[0]
try:
self.assertEqual(context.request_context.user.id, "t2_example")
self.assertEqual(context.request_context.user.roles, set())
self.assertEqual(context.request_context.user.is_logged_in, True)
self.assertEqual(context.request_context.user.loid, "t2_deadbeef")
self.assertEqual(context.request_context.user.cookie_created_ms,
100000)
self.assertEqual(context.request_context.oauth_client.id, None)
self.assertFalse(
context.request_context.oauth_client.is_type("third_party"))
self.assertEqual(context.request_context.session.id, "beefdead")
except jwt.exceptions.InvalidAlgorithmError:
raise unittest.SkipTest("cryptography is not installed")
def setUp(self):
self.itrans = TMemoryBuffer()
self.iprot = THeaderProtocol(self.itrans)
self.otrans = TMemoryBuffer()
self.oprot = THeaderProtocol(self.otrans)
self.observer = mock.Mock(spec=BaseplateObserver)
self.server_observer = mock.Mock(spec=ServerSpanObserver)
def _register_mock(context, server_span):
server_span.register(self.server_observer)
self.observer.on_server_span_created.side_effect = _register_mock
self.logger = mock.Mock(spec=logging.Logger)
self.server_context = TRpcConnectionContext(self.itrans, self.iprot,
self.oprot)
mock_filewatcher = mock.Mock(spec=FileWatcher)
mock_filewatcher.get_data.return_value = {
"secrets": {
"jwt/authentication/secret": {
"type": "simple",
"value": self.TOKEN_SECRET,
},
},
"vault": {
"token": "test",
"url": "http://vault.example.com:8200/",
}
}
secrets = store.SecretsStore("/secrets")
secrets._filewatcher = mock_filewatcher
baseplate = Baseplate()
baseplate.register(self.observer)
event_handler = BaseplateProcessorEventHandler(
self.logger,
baseplate,
auth_factory=AuthenticationContextFactory(secrets))
handler = TestHandler()
self.processor = TestService.ContextProcessor(handler)
self.processor.setEventHandler(event_handler)
def to_string(self, download_req_info):
str_req = None
try:
tMemory_b = TMemoryBuffer()
tBinaryProtocol_b = TBinaryProtocol(tMemory_b)
download_req_info.write(tBinaryProtocol_b)
str_req = tMemory_b.getvalue()
except EOFError, e:
self.log.warning("cann't write DownLoadReq to string")
def to_string(self, link_info):
str_entity = None
try:
tMemory_b = TMemoryBuffer()
tBinaryProtocol_b = TBinaryProtocol.TBinaryProtocol(tMemory_b)
link_info.write(tBinaryProtocol_b)
str_entity = tMemory_b.getvalue()
except EOFError, e:
self.log.warning("can't write LinkAttr to string")
def to_string(page_info):
str_page_info = None
try:
tMemory_b = TMemoryBuffer()
tBinaryProtocol_b = TBinaryServerProtocol(tMemory_b)
page_info.write(tBinaryProtocol_b)
str_page_info = tMemory_b.getvalue()
except EOFError, e:
pass
def write_communication_to_buffer(comm):
'''
Serialize communication to buffer (binary string) and return
buffer.
'''
transport = TMemoryBuffer()
protocol = factory.createProtocol(transport)
comm.write(protocol)
return transport.getvalue()
def to_string(log, page_info):
str_page_info = None
try:
tMemory_b = TMemoryBuffer()
tBinaryProtocol_b = TBinaryServerProtocol(tMemory_b)
page_info.write(tBinaryProtocol_b)
str_page_info = tMemory_b.getvalue()
except EOFError, e:
log.warning("cann't write DownLoadRsp to string")
def flush(self):
msg = self._writeBuffer.getvalue()
tmi = TMemoryBuffer(msg)
iprot = THeaderProtocol(tmi)
fname, mtype, seqid = iprot.readMessageBegin()
fname = fname.decode()
self._proto.schedule_timeout(fname, seqid)
self._trans.send_message(msg)
self._writeBuffer = BytesIO()
def call_processor(self, input, headers, client_type, protocol_type,
context_data):
try:
# The input string has already had the header removed, but
# the python processor will expect it to be there. In
# order to reconstitute the message with headers, we use
# the THeaderProtocol object to write into a memory
# buffer, then pass that buffer to the python processor.
write_buf = TMemoryBuffer()
trans = THeaderTransport(write_buf)
trans._THeaderTransport__client_type = client_type
trans._THeaderTransport__write_headers = headers
trans.set_protocol_id(protocol_type)
trans.write(input)
trans.flush()
prot_buf = TMemoryBuffer(write_buf.getvalue())
prot = THeaderProtocol(prot_buf, client_types=[client_type])
ctx = TCppConnectionContext(context_data)
self.processor.process(prot, prot, ctx)
# Check for empty result. If so, return an empty string
# here. This is probably a oneway request, but we can't
# reliably tell. The C++ code does basically the same
# thing.
response = prot_buf.getvalue()
if len(response) == 0:
return response
# And on the way out, we need to strip off the header,
# because the C++ code will expect to add it.
read_buf = TMemoryBuffer(response)
trans = THeaderTransport(read_buf, client_types=[client_type])
trans.readFrame(len(response))
return trans.cstringio_buf.read()
except:
# Don't let exceptions escape back into C++
traceback.print_exc()
def test_forward_compatibility_nested_nested(self):
obj = OldStructureNested()
obj.features = [{}]
obj.features[0][1] = 314
obj.features[0][2] = 271
objN = OldStructureNestedNested()
objN.field = obj
trans = TMemoryBuffer()
proto = self.createProto(trans)
objN.write(proto)
obj_new = NewStructureNestedNested()
trans = TMemoryBuffer(trans.getvalue())
proto = proto.__class__(trans)
fastproto.decode(
obj_new,
trans, [obj_new.__class__, obj_new.thrift_spec,
obj_new.isUnion()],
utf8strings=0,
protoid=self.PROTO,
forward_compatibility=True)
self.assertAlmostEqual(obj_new.field.features[0][1], 314.0)
self.assertAlmostEqual(obj_new.field.features[0][2], 271.0)
trans2 = TMemoryBuffer()
proto2 = self.createProto(trans2)
obj_new.write(proto2)
obj_new2 = NewStructureNestedNested()
trans2 = TMemoryBuffer(trans2.getvalue())
proto2 = proto2.__class__(trans2)
fastproto.decode(
obj_new2,
trans2,
[obj_new2.__class__, obj_new2.thrift_spec,
obj_new2.isUnion()],
utf8strings=0,
protoid=self.PROTO)
self.assertAlmostEqual(obj_new2.field.features[0][1], 314.0)
self.assertAlmostEqual(obj_new2.field.features[0][2], 271.0)
def decode(self, recvdata):
lenth = len(recvdata)
if lenth <= 0:
return None, None
m = TMemoryBuffer(recvdata)
p = TBinaryProtocol(m)
packet_header_len, protocol_type = self.decode_packet_header(p)
try:
h = msg_header()
h.read(p)
except EOFError:
print("Error: msg header decode failure")
return None, None
# print("package_len:")
# print(packet_header_len)
# print "msg header: "
# pprint(vars(h), indent=4)
msg_name = h.msg_name
print("recv msg name:", msg_name)
if msg_name is None:
print("big error ", binascii.hexlify(recvdata))
return 0
try:
s = h.msg_name + "_body"
if s in globals():
t = globals()[s]
else:
t = empty
body = t()
body.read(p)
except EOFError:
print("Error: msg body decode failure")
return None, None
# print("body: ")
# pprint(vars(body), indent=4, width=24)
if msg_name == "ver_resp":
self.deal_ver_resp()
if msg_name != "shake_hand_resp":
print("------------------------------------------------")
print("msg header: ")
pprint(vars(h), indent=4)
print("body: ")
pprint(vars(body), indent=4, width=24)
message = inner_message()
message.header.msg_name = str(msg_name)
message.header.src_id = self.get_sid()
message.msg_head = h
message.msg_content = body
return msg_name, message
return None, None
def do_task(self, body):
link_info = LinkAttr()
try:
tMemory_o = TMemoryBuffer(body)
tBinaryProtocol_o = TBinaryProtocol(tMemory_o)
link_info.read(tBinaryProtocol_o)
self.scheduler.schedule_task(link_info)
except EOFError, e:
self.log.warning("can't read LinkAttr from string")
return None
def to_string(self, page_info):
str_page_info = None
try:
tMemory_b = TMemoryBuffer()
tBinaryProtocol_b = TBinaryServerProtocol(tMemory_b)
page_info.write(tBinaryProtocol_b)
str_page_info = tMemory_b.getvalue()
except EOFError as e:
self._log.warning("cann't write data to string")
return str_page_info
def to_string(self, download_rsp):
str_rsq = None
try:
tMemory_b = TMemoryBuffer()
tBinaryProtocol_b = TBinaryProtocol(tMemory_b)
download_rsp.write(tBinaryProtocol_b)
str_rsq = tMemory_b.getvalue()
self.log.info('data-length is {}'.format(str(len(str_rsq))))
except EOFError, e:
self.log.warning("cann't write PageParseInfo to string")
