This repository has been archived by the owner on Dec 15, 2022. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 8
/
storage.py
133 lines (102 loc) · 4.04 KB
/
storage.py
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
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
from proto import rbperf_pb2
import struct
MAGIC_STRING = b"rbperf"
class CompactProtobufWriter:
def __init__(self, file_object):
self.file = file_object
self._write_magic_string()
self._write_storage_type()
def write_profile(self, profile):
self.file.write(profile.SerializeToString())
def _write_magic_string(self):
self.file.write(MAGIC_STRING)
def _write_storage_type(self):
self.file.write(struct.pack("<H", rbperf_pb2.StorageType.COMPACT))
class CompactProtobufReader:
def __init__(self, file_object):
self.file = file_object
self._verify_magic_string()
self._verify_storage_type()
def read_stacks(self):
self.profile = self._read_profile()
for stack in self.profile.stacktraces:
stack_trace = rbperf_pb2.StackTrace(
timestamp=stack.timestamp,
tid=stack.tid,
cpu=stack.cpu,
pid=stack.pid,
comm=stack.comm,
stack_status=stack.stack_status,
)
for frame in stack.interned_frames:
cur_frame = rbperf_pb2.Frame()
cur_frame.path = self.symbol_for_id(frame.path)
cur_frame.method = self.symbol_for_id(frame.method)
cur_frame.lineno = frame.lineno
stack_trace.frames.append(cur_frame)
yield stack_trace
def symbol_for_id(self, id_):
return self.profile.string_table[id_]
def _verify_magic_string(self):
assert self.file.read(len(MAGIC_STRING)) == MAGIC_STRING
def _verify_storage_type(self):
storage_type_bytes = self.file.read(2)
assert (
struct.unpack("<H", storage_type_bytes)[0] == rbperf_pb2.StorageType.COMPACT
)
def _read_profile(self):
profile_bytes = self.file.read()
profile = rbperf_pb2.Profile()
profile.ParseFromString(profile_bytes)
return profile
class StreamingProtobufWriter:
def __init__(self, file_object):
self.file = file_object
self._write_magic_string()
self._write_storage_type()
def write_header(self, header):
self.file.write(struct.pack("<I", header.ByteSize()))
self.file.write(header.SerializeToString())
def write_stack(self, stack):
self.file.write(struct.pack("<I", stack.ByteSize()))
self.file.write(stack.SerializeToString())
def _write_magic_string(self):
self.file.write(MAGIC_STRING)
def _write_storage_type(self):
self.file.write(struct.pack("<H", rbperf_pb2.StorageType.STREAMING))
class StreamingProtobufReader:
def __init__(self, file_object):
self.file = file_object
self._verify_magic_string()
self._verify_storage_type()
self.header = self._read_header()
def read_header(self):
return self.header
def read_stacks(self):
while True:
stack_trace_size_bytes = self.file.read(4)
if stack_trace_size_bytes == b"":
return
stack_trace_size = struct.unpack("<I", stack_trace_size_bytes)[0]
stack_trace_bytes = self.file.read(stack_trace_size)
stack_trace = rbperf_pb2.StackTrace()
stack_trace.ParseFromString(stack_trace_bytes)
yield stack_trace
def _verify_magic_string(self):
assert self.file.read(len(MAGIC_STRING)) == MAGIC_STRING
def _verify_storage_type(self):
storage_type_bytes = self.file.read(2)
assert (
struct.unpack("<H", storage_type_bytes)[0]
== rbperf_pb2.StorageType.STREAMING
)
def _read_header(self):
header_size = struct.unpack("<I", self.file.read(4))[0]
header_bytes = self.file.read(header_size)
header = rbperf_pb2.Profile()
header.ParseFromString(header_bytes)
return header