已开启
ubiobuf_reactor #152
ubiobuf_reactor #152
已开启
lipeng_2018创建于 7月6日
7 个文件变更+304-717
@@ -81,46 +81,6 @@ inline void PackRpcHeader(char* rpc_header, uint32_t meta_size, int payload_size
81 .pack32(meta_size);81 .pack32(meta_size);
82}82}
83 83 
84-#ifdef BRPC_WITH_URMA
85-static inline butil::IOBuf* SelectIOBuf(bool use_ub, butil::IOBuf* buf,
86- butil::UBIOBuf* ub_buf) {
87- return use_ub ? static_cast<butil::IOBuf*>(ub_buf) : buf;
88-}
89-#endif
90- 
91-static inline bool SerializeToIOBuf(Serializer& serializer, butil::IOBuf* out, uint32_t block_size) {
92-#ifdef BRPC_WITH_URMA
93- if (out->use_ub()) {
94- butil::UBIOBufAsZeroCopyOutputStream stream(static_cast<butil::UBIOBuf*>(out), block_size);
95- return serializer.SerializeTo(&stream);
96- }
97-#endif
98- butil::IOBufAsZeroCopyOutputStream stream(out);
99- return serializer.SerializeTo(&stream);
100-}
101- 
102-template <typename OutputStream>
103-static inline void SerializeCachedMetaToOutputStream(const RpcMeta& meta,
104- OutputStream* stream) {
105- ::google::protobuf::io::CodedOutputStream coded_out(stream);
106- meta.SerializeWithCachedSizes(&coded_out);
107- CHECK(!coded_out.HadError());
108-}
109- 
110-static inline void SerializeCachedMetaToIOBuf(const RpcMeta& meta, butil::IOBuf* out,
111- uint32_t meta_size) {
112-#ifdef BRPC_WITH_URMA
113- if (out->use_ub()) {
114- butil::UBIOBufAsZeroCopyOutputStream buf_stream(static_cast<butil::UBIOBuf*>(out),
115- meta_size);
116- SerializeCachedMetaToOutputStream(meta, &buf_stream);
117- return;
118- }
119-#endif
120- butil::IOBufAsZeroCopyOutputStream buf_stream(out);
121- SerializeCachedMetaToOutputStream(meta, &buf_stream);
122-}
123- 
124static inline int AppendRpcHeaderAndMeta(butil::IOBuf* out, const void* data,84static inline int AppendRpcHeaderAndMeta(butil::IOBuf* out, const void* data,
125 size_t size, uint32_t meta_size) {85 size_t size, uint32_t meta_size) {
126#ifdef BRPC_WITH_URMA86#ifdef BRPC_WITH_URMA
@@ -132,6 +92,11 @@ static inline int AppendRpcHeaderAndMeta(butil::IOBuf* out, const void* data,
132 return out->append(data, size);92 return out->append(data, size);
133}93}
134 94 
95+static inline butil::IOBuf* SelectIOBuf(bool use_ub, butil::IOBuf* buf,
96+ butil::UBIOBuf* ub_buf) {
97+ return use_ub ? static_cast<butil::IOBuf*>(ub_buf) : buf;
98+}
99+ 
135static void SerializeRpcHeaderAndMeta(100static void SerializeRpcHeaderAndMeta(
136 butil::IOBuf* out, const RpcMeta& meta, int payload_size) {101 butil::IOBuf* out, const RpcMeta& meta, int payload_size) {
137 const uint32_t meta_size = GetProtobufByteSize(meta);102 const uint32_t meta_size = GetProtobufByteSize(meta);
@@ -148,7 +113,24 @@ static void SerializeRpcHeaderAndMeta(
148 char header[12];113 char header[12];
149 PackRpcHeader(header, meta_size, payload_size);114 PackRpcHeader(header, meta_size, payload_size);
150 CHECK_EQ(0, AppendRpcHeaderAndMeta(out, header, sizeof(header), meta_size + payload_size));115 CHECK_EQ(0, AppendRpcHeaderAndMeta(out, header, sizeof(header), meta_size + payload_size));
151- SerializeCachedMetaToIOBuf(meta, out, meta_size + payload_size);116+#ifdef BRPC_WITH_URMA
117+ if (out->use_ub()) {
118+ butil::IOBufAsZeroCopyOutputStream buf_stream(out, meta_size + payload_size);
119+ ::google::protobuf::io::CodedOutputStream coded_out(&buf_stream);
120+ meta.SerializeWithCachedSizes(&coded_out);
121+ CHECK(!coded_out.HadError());
122+ } else {
123+ butil::IOBufAsZeroCopyOutputStream buf_stream(out);
124+ ::google::protobuf::io::CodedOutputStream coded_out(&buf_stream);
125+ meta.SerializeWithCachedSizes(&coded_out);
126+ CHECK(!coded_out.HadError());
127+ }
128+#else
129+ butil::IOBufAsZeroCopyOutputStream buf_stream(out);
130+ ::google::protobuf::io::CodedOutputStream coded_out(&buf_stream);
131+ meta.SerializeWithCachedSizes(&coded_out);
132+ CHECK(!coded_out.HadError());
133+#endif
152 }134 }
153}135}
154 136 
@@ -195,6 +177,17 @@ ParseResult ParseRpcMessage(butil::IOBuf* source, Socket* socket,
195 return MakeMessage(msg);177 return MakeMessage(msg);
196}178}
197 179 
180+static inline bool SerializeToIOBuf(Serializer& serializer, butil::IOBuf* out, uint32_t block_size) {
181+#ifdef BRPC_WITH_URMA
182+ if (out->use_ub()) {
183+ butil::IOBufAsZeroCopyOutputStream stream(out, block_size);
184+ return serializer.SerializeTo(&stream);
185+ }
186+#endif
187+ butil::IOBufAsZeroCopyOutputStream stream(out);
188+ return serializer.SerializeTo(&stream);
189+}
190+ 
198bool SerializeRpcMessage(const google::protobuf::Message& message,191bool SerializeRpcMessage(const google::protobuf::Message& message,
199 Controller& cntl, ContentType content_type,192 Controller& cntl, ContentType content_type,
200 CompressType compress_type, ChecksumType checksum_type,193 CompressType compress_type, ChecksumType checksum_type,
@@ -375,14 +368,10 @@ void SendRpcResponse(int64_t correlation_id, Controller* cntl,
375 return;368 return;
376 }369 }
377 bool append_body = false;370 bool append_body = false;
378- butil::IOBuf res_body;
379-#ifdef BRPC_WITH_URMA
380 const bool use_ub = sock->use_ub();371 const bool use_ub = sock->use_ub();
372+ butil::IOBuf res_body;
381 butil::UBIOBuf ub_res_body;373 butil::UBIOBuf ub_res_body;
382 butil::IOBuf* res_body_ptr = SelectIOBuf(use_ub, &res_body, &ub_res_body);374 butil::IOBuf* res_body_ptr = SelectIOBuf(use_ub, &res_body, &ub_res_body);
383-#else
384- butil::IOBuf* res_body_ptr = &res_body;
385-#endif
386 // `res' can be NULL here, in which case we don't serialize it375 // `res' can be NULL here, in which case we don't serialize it
387 // If user calls `SetFailed' on Controller, we don't serialize376 // If user calls `SetFailed' on Controller, we don't serialize
388 // response either377 // response either
@@ -454,12 +443,8 @@ void SendRpcResponse(int64_t correlation_id, Controller* cntl,
454 }443 }
455 444 
456 butil::IOBuf res_buf;445 butil::IOBuf res_buf;
457-#ifdef BRPC_WITH_URMA
458 butil::UBIOBuf ub_res_buf;446 butil::UBIOBuf ub_res_buf;
459 butil::IOBuf* res_buf_ptr = SelectIOBuf(use_ub, &res_buf, &ub_res_buf);447 butil::IOBuf* res_buf_ptr = SelectIOBuf(use_ub, &res_buf, &ub_res_buf);
460-#else
461- butil::IOBuf* res_buf_ptr = &res_buf;
462-#endif
463 SerializeRpcHeaderAndMeta(res_buf_ptr, meta, res_size + attached_size);448 SerializeRpcHeaderAndMeta(res_buf_ptr, meta, res_size + attached_size);
464 if (append_body) {449 if (append_body) {
465 res_buf_ptr->append(res_body_ptr->movable());450 res_buf_ptr->append(res_body_ptr->movable());
@@ -976,12 +961,8 @@ bool VerifyRpcRequest(const InputMessageBase* msg_base) {
976 response_meta.mutable_response()->mutable_error_text()->append(user_error_text);961 response_meta.mutable_response()->mutable_error_text()->append(user_error_text);
977 }962 }
978 butil::IOBuf res_buf;963 butil::IOBuf res_buf;
979-#ifdef BRPC_WITH_URMA
980 butil::UBIOBuf ub_res_buf;964 butil::UBIOBuf ub_res_buf;
981 butil::IOBuf* res_buf_ptr = SelectIOBuf(socket->use_ub(), &res_buf, &ub_res_buf);965 butil::IOBuf* res_buf_ptr = SelectIOBuf(socket->use_ub(), &res_buf, &ub_res_buf);
982-#else
983- butil::IOBuf* res_buf_ptr = &res_buf;
984-#endif
985 SerializeRpcHeaderAndMeta(res_buf_ptr, response_meta, 0);966 SerializeRpcHeaderAndMeta(res_buf_ptr, response_meta, 0);
986 Socket::WriteOptions opt;967 Socket::WriteOptions opt;
987 opt.ignore_eovercrowded = true;968 opt.ignore_eovercrowded = true;
@@ -17,7 +17,6 @@
17 17 
18 18 
19#include "butil/logging.h"19#include "butil/logging.h"
20-#include "butil/ub/ubiobuf.h" // butil::UBIOBufAsSnappySink
21#include "butil/third_party/snappy/snappy.h"20#include "butil/third_party/snappy/snappy.h"
22#include "brpc/policy/snappy_compress.h"21#include "brpc/policy/snappy_compress.h"
23#include "brpc/protocol.h"22#include "brpc/protocol.h"
@@ -26,22 +25,6 @@
26namespace brpc {25namespace brpc {
27namespace policy {26namespace policy {
28 27 
29-namespace {
30- 
31-bool SnappyCompressToIOBuf(const butil::IOBuf& in, butil::IOBuf* out) {
32- butil::IOBufAsSnappySource source(in);
33-#if BRPC_WITH_URMA
34- if (out->use_ub()) {
35- butil::UBIOBufAsSnappySink sink(*static_cast<butil::UBIOBuf*>(out));
36- return butil::snappy::Compress(&source, &sink);
37- }
38-#endif
39- butil::IOBufAsSnappySink sink(*out);
40- return butil::snappy::Compress(&source, &sink);
41-}
42- 
43-} // namespace
44- 
45bool SnappyCompress(const google::protobuf::Message& msg, butil::IOBuf* buf) {28bool SnappyCompress(const google::protobuf::Message& msg, butil::IOBuf* buf) {
46 butil::IOBuf serialized_pb;29 butil::IOBuf serialized_pb;
47 butil::IOBufAsZeroCopyOutputStream wrapper(&serialized_pb);30 butil::IOBufAsZeroCopyOutputStream wrapper(&serialized_pb);
@@ -87,7 +70,9 @@ bool SnappyDecompress(const butil::IOBuf& data, google::protobuf::Message* msg)
87}70}
88 71 
89bool SnappyCompress(const butil::IOBuf& in, butil::IOBuf* out) {72bool SnappyCompress(const butil::IOBuf& in, butil::IOBuf* out) {
90- return SnappyCompressToIOBuf(in, out);73+ butil::IOBufAsSnappySource source(in);
74+ butil::IOBufAsSnappySink sink(*out);
75+ return butil::snappy::Compress(&source, &sink);
91}76}
92 77 
93bool SnappyDecompress(const butil::IOBuf& in, butil::IOBuf* out) {78bool SnappyDecompress(const butil::IOBuf& in, butil::IOBuf* out) {
@@ -259,8 +259,6 @@ inline IOBuf::Block* create_block_aligned(size_t block_size, size_t alignment) {
259 259 
260static __thread TLSData g_tls_data = { NULL, 0, false };260static __thread TLSData g_tls_data = { NULL, 0, false };
261 261 
262-// Used in release_tls_block()
263-TLSData* get_g_tls_data() { return &g_tls_data; }
264// Used in UT262// Used in UT
265IOBuf::Block* get_tls_block_head() { return g_tls_data.block_head; }263IOBuf::Block* get_tls_block_head() { return g_tls_data.block_head; }
266int get_tls_block_count() { return g_tls_data.num_blocks; }264int get_tls_block_count() { return g_tls_data.num_blocks; }
@@ -278,113 +276,34 @@ void dec_g_num_hit_tls_threshold() {
278 g_num_hit_tls_threshold.fetch_sub(1, butil::memory_order_relaxed);276 g_num_hit_tls_threshold.fetch_sub(1, butil::memory_order_relaxed);
279}277}
280 278 
279+TLSData* IOBufTLSBlockPolicy::tls_data() {
280+ return &g_tls_data;
281+}
282+ 
283+void IOBufTLSBlockPolicy::inc_hit_tls_threshold(size_t n) {
284+ g_num_hit_tls_threshold.fetch_add(n, butil::memory_order_relaxed);
285+}
286+ 
281// Called in UT.287// Called in UT.
282void remove_tls_block_chain() {288void remove_tls_block_chain() {
283- TLSData& tls_data = g_tls_data;289+ detail::remove_tls_block_chain<IOBufTLSBlockPolicy>();
284- IOBuf::Block* b = tls_data.block_head;
285- if (!b) {
286- return;
287- }
288- tls_data.block_head = NULL;
289- int n = 0;
290- do {
291- IOBuf::Block* const saved_next = b->u.portal_next;
292- b->dec_ref();
293- b = saved_next;
294- ++n;
295- } while (b);
296- CHECK_EQ(n, tls_data.num_blocks);
297- tls_data.num_blocks = 0;
298}290}
299 291 
300// Get a (non-full) block from TLS.292// Get a (non-full) block from TLS.
301// Notice that the block is not removed from TLS.293// Notice that the block is not removed from TLS.
302IOBuf::Block* share_tls_block() {294IOBuf::Block* share_tls_block() {
303- TLSData& tls_data = g_tls_data;295+ return detail::share_tls_block<IOBufTLSBlockPolicy>();
304- IOBuf::Block* const b = tls_data.block_head;
305- if (b != NULL && !b->full()) {
306- return b;
307- }
308- IOBuf::Block* new_block = NULL;
309- if (b) {
310- new_block = b;
311- while (new_block && new_block->full()) {
312- IOBuf::Block* const saved_next = new_block->u.portal_next;
313- new_block->dec_ref();
314- --tls_data.num_blocks;
315- new_block = saved_next;
316- }
317- } else if (!tls_data.registered) {
318- tls_data.registered = true;
319- // Only register atexit at the first time
320- butil::thread_atexit(remove_tls_block_chain);
321- }
322- if (!new_block) {
323- new_block = create_block(); // may be NULL
324- if (new_block) {
325- ++tls_data.num_blocks;
326- }
327- }
328- tls_data.block_head = new_block;
329- return new_block;
330}296}
331 297 
332// Return chained blocks to TLS.298// Return chained blocks to TLS.
333// NOTE: b MUST be non-NULL and all blocks linked SHOULD not be full.299// NOTE: b MUST be non-NULL and all blocks linked SHOULD not be full.
334void release_tls_block_chain(IOBuf::Block* b) {300void release_tls_block_chain(IOBuf::Block* b) {
335- TLSData& tls_data = g_tls_data;301+ detail::release_tls_block_chain<IOBufTLSBlockPolicy>(b);
336- size_t n = 0;
337- if (tls_data.num_blocks >= max_blocks_per_thread()) {
338- do {
339- ++n;
340- IOBuf::Block* const saved_next = b->u.portal_next;
341- b->dec_ref();
342- b = saved_next;
343- } while (b);
344- g_num_hit_tls_threshold.fetch_add(n, butil::memory_order_relaxed);
345- return;
346- }
347- IOBuf::Block* first_b = b;
348- IOBuf::Block* last_b = NULL;
349- do {
350- ++n;
351- CHECK(!b->full());
352- if (b->u.portal_next == NULL) {
353- last_b = b;
354- break;
355- }
356- b = b->u.portal_next;
357- } while (true);
358- last_b->u.portal_next = tls_data.block_head;
359- tls_data.block_head = first_b;
360- tls_data.num_blocks += n;
361- if (!tls_data.registered) {
362- tls_data.registered = true;
363- butil::thread_atexit(remove_tls_block_chain);
364- }
365}302}
366 303 
367// Get and remove one (non-full) block from TLS. If TLS is empty, create one.304// Get and remove one (non-full) block from TLS. If TLS is empty, create one.
368IOBuf::Block* acquire_tls_block() {305IOBuf::Block* acquire_tls_block() {
369- TLSData& tls_data = g_tls_data;306+ return detail::acquire_tls_block<IOBufTLSBlockPolicy>();
370- IOBuf::Block* b = tls_data.block_head;
371- if (!b) {
372- return create_block();
373- }
374- while (b->full()) {
375- IOBuf::Block* const saved_next = b->u.portal_next;
376- b->dec_ref();
377- tls_data.block_head = saved_next;
378- --tls_data.num_blocks;
379- b = saved_next;
380- if (!b) {
381- return create_block();
382- }
383- }
384- tls_data.block_head = b->u.portal_next;
385- --tls_data.num_blocks;
386- b->u.portal_next = NULL;
387- return b;
388}307}
389 308 
390inline IOBuf::BlockRef* acquire_blockref_array(size_t cap) {309inline IOBuf::BlockRef* acquire_blockref_array(size_t cap) {
@@ -442,12 +361,12 @@ void IOBuf::operator=(const IOBuf& rhs) {
442 return;361 return;
443 }362 }
444 if (!rhs._small() && !_small() && _bv.cap_mask == rhs._bv.cap_mask) {363 if (!rhs._small() && !_small() && _bv.cap_mask == rhs._bv.cap_mask) {
445- // Reuse array of refs364+ // Reuse array of refs.
446 // Remove references to previous blocks.365 // Remove references to previous blocks.
447 for (size_t i = 0; i < _bv.nref; ++i) {366 for (size_t i = 0; i < _bv.nref; ++i) {
448 _bv.ref_at(i).block->dec_ref();367 _bv.ref_at(i).block->dec_ref();
449 }368 }
450- // References blocks in rhs.369+ // Reference blocks in rhs.
451 _bv.start = 0;370 _bv.start = 0;
452 _bv.nref = rhs._bv.nref;371 _bv.nref = rhs._bv.nref;
453 _bv.nbytes = rhs._bv.nbytes;372 _bv.nbytes = rhs._bv.nbytes;
@@ -456,8 +375,8 @@ void IOBuf::operator=(const IOBuf& rhs) {
456 _bv.refs[i].block->inc_ref();375 _bv.refs[i].block->inc_ref();
457 }376 }
458 } else {377 } else {
459- this->~IOBuf();378+ clear();
460- new (this) IOBuf(rhs);379+ append(rhs);
461 }380 }
462}381}
463 382 
@@ -465,6 +384,22 @@ bool IOBuf::use_ub() const {
465 return false;384 return false;
466}385}
467 386 
387+IOBuf::Block* IOBuf::_share_tls_block() {
388+ return iobuf::share_tls_block();
389+}
390+ 
391+IOBuf::Block* IOBuf::_create_block(uint32_t block_size) {
392+ return iobuf::create_block(block_size);
393+}
394+ 
395+IOBuf::Block* IOBuf::_acquire_tls_block() {
396+ return iobuf::acquire_tls_block();
397+}
398+ 
399+void IOBuf::_release_tls_block(IOBuf::Block* block) {
400+ iobuf::release_tls_block(block);
401+}
402+ 
468template <bool MOVE>403template <bool MOVE>
469void IOBuf::_push_or_move_back_ref_to_smallview(const BlockRef& r) {404void IOBuf::_push_or_move_back_ref_to_smallview(const BlockRef& r) {
470 BlockRef* const refs = _sv.refs;405 BlockRef* const refs = _sv.refs;
@@ -1089,7 +1024,7 @@ void IOBuf::append(const Movable& movable_other) {
1089}1024}
1090 1025 
1091int IOBuf::push_back(char c) {1026int IOBuf::push_back(char c) {
1092- IOBuf::Block* b = iobuf::share_tls_block();1027+ IOBuf::Block* b = _share_tls_block();
1093 if (BAIDU_UNLIKELY(!b)) {1028 if (BAIDU_UNLIKELY(!b)) {
1094 return -1;1029 return -1;
1095 }1030 }
@@ -1116,7 +1051,7 @@ int IOBuf::append(void const* data, size_t count) {
1116 }1051 }
1117 size_t total_nc = 0;1052 size_t total_nc = 0;
1118 while (total_nc < count) { // excluded count == 01053 while (total_nc < count) { // excluded count == 0
1119- IOBuf::Block* b = iobuf::share_tls_block();1054+ IOBuf::Block* b = _share_tls_block();
1120 if (BAIDU_UNLIKELY(!b)) {1055 if (BAIDU_UNLIKELY(!b)) {
1121 return -1;1056 return -1;
1122 }1057 }
@@ -1134,7 +1069,7 @@ int IOBuf::append(void const* data, size_t count) {
1134int IOBuf::appendv(const const_iovec* vec, size_t n) {1069int IOBuf::appendv(const const_iovec* vec, size_t n) {
1135 size_t offset = 0;1070 size_t offset = 0;
1136 for (size_t i = 0; i < n;) {1071 for (size_t i = 0; i < n;) {
1137- IOBuf::Block* b = iobuf::share_tls_block();1072+ IOBuf::Block* b = _share_tls_block();
1138 if (BAIDU_UNLIKELY(!b)) {1073 if (BAIDU_UNLIKELY(!b)) {
1139 return -1;1074 return -1;
1140 }1075 }
@@ -1203,7 +1138,7 @@ int IOBuf::resize(size_t n, char c) {
1203 const size_t count = n - saved_len;1138 const size_t count = n - saved_len;
1204 size_t total_nc = 0;1139 size_t total_nc = 0;
1205 while (total_nc < count) { // excluded count == 01140 while (total_nc < count) { // excluded count == 0
1206- IOBuf::Block* b = iobuf::share_tls_block();1141+ IOBuf::Block* b = _share_tls_block();
1207 if (BAIDU_UNLIKELY(!b)) {1142 if (BAIDU_UNLIKELY(!b)) {
1208 return -1;1143 return -1;
1209 }1144 }
@@ -1252,7 +1187,7 @@ IOBuf::Area IOBuf::reserve(size_t count) {
1252 IOBuf::Area result = INVALID_AREA;1187 IOBuf::Area result = INVALID_AREA;
1253 size_t total_nc = 0;1188 size_t total_nc = 0;
1254 while (total_nc < count) { // excluded count == 01189 while (total_nc < count) { // excluded count == 0
1255- IOBuf::Block* b = iobuf::share_tls_block();1190+ IOBuf::Block* b = _share_tls_block();
1256 if (BAIDU_UNLIKELY(!b)) {1191 if (BAIDU_UNLIKELY(!b)) {
1257 return INVALID_AREA;1192 return INVALID_AREA;
1258 }1193 }
@@ -1934,7 +1869,7 @@ IOBufAsZeroCopyOutputStream::IOBufAsZeroCopyOutputStream(
1934 , _cur_block(NULL)1869 , _cur_block(NULL)
1935 , _byte_count(0) {1870 , _byte_count(0) {
1936 1871
1937- if (_block_size <= offsetof(IOBuf::Block, data)) {1872+ if (!_buf->use_ub() && _block_size <= offsetof(IOBuf::Block, data)) {
1938 throw std::invalid_argument("block_size is too small");1873 throw std::invalid_argument("block_size is too small");
1939 }1874 }
1940}1875}
@@ -1947,9 +1882,9 @@ bool IOBufAsZeroCopyOutputStream::Next(void** data, int* size) {
1947 if (_cur_block == NULL || _cur_block->full()) {1882 if (_cur_block == NULL || _cur_block->full()) {
1948 _release_block();1883 _release_block();
1949 if (_block_size > 0) {1884 if (_block_size > 0) {
1950- _cur_block = iobuf::create_block(_block_size);1885+ _cur_block = _buf->_create_block(_block_size);
1951 } else {1886 } else {
1952- _cur_block = iobuf::acquire_tls_block();1887+ _cur_block = _buf->_acquire_tls_block();
1953 }1888 }
1954 if (_cur_block == NULL) {1889 if (_cur_block == NULL) {
1955 return false;1890 return false;
@@ -2022,8 +1957,8 @@ void IOBufAsZeroCopyOutputStream::BackUp(int count) {
2022 // ParseFromZeroCopyStream(&wrapper, ...); // Calls BackUp1957 // ParseFromZeroCopyStream(&wrapper, ...); // Calls BackUp
2023 // IOBuf buf;1958 // IOBuf buf;
2024 // buf.append("foobar"); // can reuse the TLS block.1959 // buf.append("foobar"); // can reuse the TLS block.
2025- if (_block_size == 0) {1960+ if (_block_size == 0 || _buf->use_ub()) {
2026- iobuf::release_tls_block(_cur_block);1961+ _buf->_release_tls_block(_cur_block);
2027 _cur_block = NULL;1962 _cur_block = NULL;
2028 }1963 }
2029 return;1964 return;
@@ -2045,12 +1980,12 @@ int64_t IOBufAsZeroCopyOutputStream::ByteCount() const {
2045}1980}
2046 1981 
2047void IOBufAsZeroCopyOutputStream::_release_block() {1982void IOBufAsZeroCopyOutputStream::_release_block() {
2048- if (_block_size > 0) {1983+ if (_block_size > 0 && !_buf->use_ub()) {
2049 if (_cur_block) {1984 if (_cur_block) {
2050 _cur_block->dec_ref();1985 _cur_block->dec_ref();
2051 }1986 }
2052 } else {1987 } else {
2053- iobuf::release_tls_block(_cur_block);1988+ _buf->_release_tls_block(_cur_block);
2054 }1989 }
2055 _cur_block = NULL;1990 _cur_block = NULL;
2056}1991}
@@ -2073,7 +2008,7 @@ void IOBufAsSnappySink::Append(const char* bytes, size_t n) {
2073char* IOBufAsSnappySink::GetAppendBuffer(size_t length, char* scratch) {2008char* IOBufAsSnappySink::GetAppendBuffer(size_t length, char* scratch) {
2074 // TODO: butil::IOBuf supports dynamic sized blocks.2009 // TODO: butil::IOBuf supports dynamic sized blocks.
2075 if (length <= 8000/*just a hint*/) {2010 if (length <= 8000/*just a hint*/) {
2076- if (_buf_stream.Next(reinterpret_cast<void**>(&_cur_buf), &_cur_len)) { 2011+ if (_buf_stream.Next(reinterpret_cast<void**>(&_cur_buf), &_cur_len)) {
2077 if (_cur_len >= static_cast<int>(length)) {2012 if (_cur_len >= static_cast<int>(length)) {
2078 return _cur_buf;2013 return _cur_buf;
2079 } else {2014 } else {
@@ -58,7 +58,6 @@ struct ssl_st;
58 58 
59namespace butil {59namespace butil {
60 60 
61-class UBIOBufAsZeroCopyOutputStream;
62class UBIOBuf;61class UBIOBuf;
63 62 
64// IOBuf is a non-continuous buffer that can be cut and combined w/o copying63// IOBuf is a non-continuous buffer that can be cut and combined w/o copying
@@ -71,7 +70,6 @@ class UBIOBuf;
71class IOBuf {70class IOBuf {
72friend class IOBufAsZeroCopyInputStream;71friend class IOBufAsZeroCopyInputStream;
73friend class IOBufAsZeroCopyOutputStream;72friend class IOBufAsZeroCopyOutputStream;
74-friend class UBIOBufAsZeroCopyOutputStream;
75friend class IOBufBytesIterator;73friend class IOBufBytesIterator;
76friend class IOBufCutter;74friend class IOBufCutter;
77friend class SingleIOBuf;75friend class SingleIOBuf;
@@ -128,10 +126,10 @@ public:
128 IOBuf(const IOBuf&);126 IOBuf(const IOBuf&);
129 IOBuf(const Movable&);127 IOBuf(const Movable&);
130 virtual ~IOBuf() { clear(); }128 virtual ~IOBuf() { clear(); }
131- virtual void operator=(const IOBuf&);129+ void operator=(const IOBuf&);
132- virtual void operator=(const Movable&);130+ void operator=(const Movable&);
133- virtual void operator=(const char*);131+ void operator=(const char*);
134- virtual void operator=(const std::string&);132+ void operator=(const std::string&);
135 133 
136 // Exchange internal fields with another IOBuf.134 // Exchange internal fields with another IOBuf.
137 void swap(IOBuf&);135 void swap(IOBuf&);
@@ -215,9 +213,9 @@ public:
215 213 
216 // Append another IOBuf to back side, payload of the IOBuf is shared214 // Append another IOBuf to back side, payload of the IOBuf is shared
217 // rather than copied.215 // rather than copied.
218- virtual void append(const IOBuf& other);216+ void append(const IOBuf& other);
219 // Append content of `other' to self and clear `other'.217 // Append content of `other' to self and clear `other'.
220- virtual void append(const Movable& other);218+ void append(const Movable& other);
221 219 
222 // ===================================================================220 // ===================================================================
223 // Following push_back()/append() are just implemented for convenience221 // Following push_back()/append() are just implemented for convenience
@@ -229,11 +227,11 @@ public:
229 227
230 // Append a character to back side. (with copying)228 // Append a character to back side. (with copying)
231 // Returns 0 on success, -1 otherwise.229 // Returns 0 on success, -1 otherwise.
232- virtual int push_back(char c);230+ int push_back(char c);
233 231
234 // Append `data' with `count' bytes to back side. (with copying)232 // Append `data' with `count' bytes to back side. (with copying)
235 // Returns 0 on success(include count == 0), -1 otherwise.233 // Returns 0 on success(include count == 0), -1 otherwise.
236- virtual int append(void const* data, size_t count);234+ int append(void const* data, size_t count);
237 235 
238 // Append multiple data to back side in one call, faster than appending236 // Append multiple data to back side in one call, faster than appending
239 // one by one separately.237 // one by one separately.
@@ -243,30 +241,30 @@ public:
243 // { data2, len2 },241 // { data2, len2 },
244 // { data3, len3 } };242 // { data3, len3 } };
245 // foo.appendv(vec, arraysize(vec));243 // foo.appendv(vec, arraysize(vec));
246- virtual int appendv(const const_iovec vec[], size_t n);244+ int appendv(const const_iovec vec[], size_t n);
247- virtual int appendv(const iovec* vec, size_t n)245+ int appendv(const iovec* vec, size_t n)
248 { return appendv((const const_iovec*)vec, n); }246 { return appendv((const const_iovec*)vec, n); }
249 247 
250 // Append a c-style string to back side. (with copying)248 // Append a c-style string to back side. (with copying)
251 // Returns 0 on success, -1 otherwise.249 // Returns 0 on success, -1 otherwise.
252 // NOTE: Returns 0 when `s' is empty.250 // NOTE: Returns 0 when `s' is empty.
253- virtual int append(char const* s);251+ int append(char const* s);
254 252 
255 // Append a std::string to back side. (with copying)253 // Append a std::string to back side. (with copying)
256 // Returns 0 on success, -1 otherwise.254 // Returns 0 on success, -1 otherwise.
257 // NOTE: Returns 0 when `s' is empty.255 // NOTE: Returns 0 when `s' is empty.
258- virtual int append(const std::string& s);256+ int append(const std::string& s);
259 257 
260 // Append the user-data to back side WITHOUT copying.258 // Append the user-data to back side WITHOUT copying.
261 // The user-data can be split and shared by smaller IOBufs and will be259 // The user-data can be split and shared by smaller IOBufs and will be
262 // deleted using the deleter func when no IOBuf references it anymore.260 // deleted using the deleter func when no IOBuf references it anymore.
263 // UBIOBuf can hold user-data blocks; normalize() copies them into UB blocks.261 // UBIOBuf can hold user-data blocks; normalize() copies them into UB blocks.
264- virtual int append_user_data(void* data, size_t size, std::function<void(void*)> deleter);262+ int append_user_data(void* data, size_t size, std::function<void(void*)> deleter);
265 263 
266 // Append the user-data to back side WITHOUT copying.264 // Append the user-data to back side WITHOUT copying.
267 // The meta is associated with this piece of user-data.265 // The meta is associated with this piece of user-data.
268 // UBIOBuf can hold user-data blocks; normalize() copies them into UB blocks.266 // UBIOBuf can hold user-data blocks; normalize() copies them into UB blocks.
269- virtual int append_user_data_with_meta(void* data, size_t size, std::function<void(void*)> deleter, uint64_t meta);267+ int append_user_data_with_meta(void* data, size_t size, std::function<void(void*)> deleter, uint64_t meta);
270 268 
271 // Get the data meta of the first byte in this IOBuf.269 // Get the data meta of the first byte in this IOBuf.
272 // The meta is specified with append_user_data_with_meta before.270 // The meta is specified with append_user_data_with_meta before.
@@ -280,8 +278,8 @@ public:
280 // as many |c| as needed to reach a size of n. If c is not specified,278 // as many |c| as needed to reach a size of n. If c is not specified,
281 // null-character would be appended.279 // null-character would be appended.
282 // Returns 0 on success, -1 otherwise.280 // Returns 0 on success, -1 otherwise.
283- virtual int resize(size_t n) { return resize(n, '\0'); }281+ int resize(size_t n) { return resize(n, '\0'); }
284- virtual int resize(size_t n, char c);282+ int resize(size_t n, char c);
285 283 
286 // Reserve `n' uninitialized bytes at back-side.284 // Reserve `n' uninitialized bytes at back-side.
287 // Returns an object representing the reserved area, INVALID_AREA on failure.285 // Returns an object representing the reserved area, INVALID_AREA on failure.
@@ -384,6 +382,11 @@ public:
384 virtual bool use_ub() const;382 virtual bool use_ub() const;
385 383 
386protected:384protected:
385+ virtual Block* _share_tls_block();
386+ virtual Block* _create_block(uint32_t block_size);
387+ virtual Block* _acquire_tls_block();
388+ virtual void _release_tls_block(Block* block);
389+ 
387 int _cut_by_char(IOBuf* out, char);390 int _cut_by_char(IOBuf* out, char);
388 int _cut_by_delim(IOBuf* out, char const* dbegin, size_t ndelim);391 int _cut_by_delim(IOBuf* out, char const* dbegin, size_t ndelim);
389 392 
@@ -609,7 +612,7 @@ class IOBufAsZeroCopyOutputStream
609public:612public:
610 explicit IOBufAsZeroCopyOutputStream(IOBuf*);613 explicit IOBufAsZeroCopyOutputStream(IOBuf*);
611 IOBufAsZeroCopyOutputStream(IOBuf*, uint32_t block_size);614 IOBufAsZeroCopyOutputStream(IOBuf*, uint32_t block_size);
612- ~IOBufAsZeroCopyOutputStream();615+ virtual ~IOBufAsZeroCopyOutputStream();
613 616 
614 bool Next(void** data, int* size) override;617 bool Next(void** data, int* size) override;
615 void BackUp(int count) override; // `count' can be as long as ByteCount()618 void BackUp(int count) override; // `count' can be as long as ByteCount()
@@ -657,7 +660,7 @@ public:
657 660
658 // Returns a writable buffer of the specified length for appending.661 // Returns a writable buffer of the specified length for appending.
659 char* GetAppendBuffer(size_t length, char* scratch) override;662 char* GetAppendBuffer(size_t length, char* scratch) override;
660- 663+ 
661private:664private:
662 char* _cur_buf;665 char* _cur_buf;
663 int _cur_len;666 int _cur_len;
@@ -668,34 +668,169 @@ inline int max_blocks_per_thread() {
668 return IsIOBufProfilerEnabled() ? 0 : MAX_BLOCKS_PER_THREAD;668 return IsIOBufProfilerEnabled() ? 0 : MAX_BLOCKS_PER_THREAD;
669}669}
670 670 
671-TLSData* get_g_tls_data();
672void remove_tls_block_chain();671void remove_tls_block_chain();
673 672 
674IOBuf::Block* acquire_tls_block();673IOBuf::Block* acquire_tls_block();
674+inline IOBuf::Block* create_block();
675 675 
676-// Return one block to TLS.676+struct IOBufTLSBlockPolicy {
677+ static TLSData* tls_data();
678+ 
679+ static int max_blocks_per_thread() {
680+ return iobuf::max_blocks_per_thread();
681+ }
682+ 
683+ static IOBuf::Block* create_block() {
684+ return iobuf::create_block();
685+ }
686+ 
687+ static void remove_tls_block_chain() {
688+ iobuf::remove_tls_block_chain();
689+ }
690+ 
691+ static void inc_hit_tls_threshold(size_t n);
692+ 
693+ static void check_block(IOBuf::Block*) {}
694+};
695+ 
696+namespace detail {
697+ 
698+template <typename Policy>
677inline void release_tls_block(IOBuf::Block* b) {699inline void release_tls_block(IOBuf::Block* b) {
678 if (!b) {700 if (!b) {
679 return;701 return;
680 }702 }
681- TLSData *tls_data = get_g_tls_data();703+ TLSData* tls_data = Policy::tls_data();
682 if (b->full()) {704 if (b->full()) {
683 b->dec_ref();705 b->dec_ref();
684- } else if (tls_data->num_blocks >= max_blocks_per_thread()) {706+ } else if (tls_data->num_blocks >= Policy::max_blocks_per_thread()) {
685 b->dec_ref();707 b->dec_ref();
686- // g_num_hit_tls_threshold.fetch_add(1, butil::memory_order_relaxed);708+ Policy::inc_hit_tls_threshold(1);
687- inc_g_num_hit_tls_threshold();
688 } else {709 } else {
689 b->u.portal_next = tls_data->block_head;710 b->u.portal_next = tls_data->block_head;
690 tls_data->block_head = b;711 tls_data->block_head = b;
691 ++tls_data->num_blocks;712 ++tls_data->num_blocks;
692 if (!tls_data->registered) {713 if (!tls_data->registered) {
693 tls_data->registered = true;714 tls_data->registered = true;
694- butil::thread_atexit(remove_tls_block_chain);715+ butil::thread_atexit(Policy::remove_tls_block_chain);
695 }716 }
696 }717 }
697}718}
698 719 
720+template <typename Policy>
721+inline void remove_tls_block_chain() {
722+ TLSData* tls_data = Policy::tls_data();
723+ IOBuf::Block* b = tls_data->block_head;
724+ if (!b) {
725+ return;
726+ }
727+ tls_data->block_head = NULL;
728+ int n = 0;
729+ do {
730+ IOBuf::Block* const saved_next = b->u.portal_next;
731+ b->dec_ref();
732+ b = saved_next;
733+ ++n;
734+ } while (b);
735+ CHECK_EQ(n, tls_data->num_blocks);
736+ tls_data->num_blocks = 0;
737+}
738+ 
739+template <typename Policy>
740+inline IOBuf::Block* share_tls_block() {
741+ TLSData* tls_data = Policy::tls_data();
742+ IOBuf::Block* const b = tls_data->block_head;
743+ if (b != NULL && !b->full()) {
744+ return b;
745+ }
746+ IOBuf::Block* new_block = NULL;
747+ if (b) {
748+ new_block = b;
749+ while (new_block && new_block->full()) {
750+ IOBuf::Block* const saved_next = new_block->u.portal_next;
751+ new_block->dec_ref();
752+ --tls_data->num_blocks;
753+ new_block = saved_next;
754+ }
755+ } else if (!tls_data->registered) {
756+ tls_data->registered = true;
757+ butil::thread_atexit(Policy::remove_tls_block_chain);
758+ }
759+ if (!new_block) {
760+ new_block = Policy::create_block(); // may be NULL
761+ if (new_block) {
762+ ++tls_data->num_blocks;
763+ }
764+ }
765+ tls_data->block_head = new_block;
766+ return new_block;
767+}
768+ 
769+template <typename Policy>
770+inline void release_tls_block_chain(IOBuf::Block* b) {
771+ TLSData* tls_data = Policy::tls_data();
772+ size_t n = 0;
773+ if (tls_data->num_blocks >= Policy::max_blocks_per_thread()) {
774+ do {
775+ ++n;
776+ IOBuf::Block* const saved_next = b->u.portal_next;
777+ b->dec_ref();
778+ b = saved_next;
779+ } while (b);
780+ Policy::inc_hit_tls_threshold(n);
781+ return;
782+ }
783+ IOBuf::Block* first_b = b;
784+ IOBuf::Block* last_b = NULL;
785+ do {
786+ ++n;
787+ CHECK(!b->full());
788+ Policy::check_block(b);
789+ if (b->u.portal_next == NULL) {
790+ last_b = b;
791+ break;
792+ }
793+ b = b->u.portal_next;
794+ } while (true);
795+ last_b->u.portal_next = tls_data->block_head;
796+ tls_data->block_head = first_b;
797+ tls_data->num_blocks += n;
798+ if (!tls_data->registered) {
799+ tls_data->registered = true;
800+ butil::thread_atexit(Policy::remove_tls_block_chain);
801+ }
802+}
803+ 
804+template <typename Policy>
805+inline IOBuf::Block* acquire_tls_block() {
806+ TLSData* tls_data = Policy::tls_data();
807+ IOBuf::Block* b = tls_data->block_head;
808+ if (!b) {
809+ return Policy::create_block();
810+ }
811+ while (b->full()) {
812+ IOBuf::Block* const saved_next = b->u.portal_next;
813+ b->dec_ref();
814+ tls_data->block_head = saved_next;
815+ --tls_data->num_blocks;
816+ b = saved_next;
817+ if (!b) {
818+ return Policy::create_block();
819+ }
820+ }
821+ tls_data->block_head = b->u.portal_next;
822+ --tls_data->num_blocks;
823+ b->u.portal_next = NULL;
824+ return b;
825+}
826+ 
827+} // namespace detail
828+ 
829+// Return one block to TLS.
830+inline void release_tls_block(IOBuf::Block* b) {
831+ detail::release_tls_block<IOBufTLSBlockPolicy>(b);
832+}
833+ 
699inline IOBuf::Block* create_block(const size_t block_size) {834inline IOBuf::Block* create_block(const size_t block_size) {
700 if (block_size > 0xFFFFFFFFULL) {835 if (block_size > 0xFFFFFFFFULL) {
701 LOG(FATAL) << "block_size=" << block_size << " is too large";836 LOG(FATAL) << "block_size=" << block_size << " is too large";
@@ -23,6 +23,7 @@
23#include "butil/ub/ubiobuf.h"23#include "butil/ub/ubiobuf.h"
24#include "butil/fd_guard.h" // butil::fd_guard24#include "butil/fd_guard.h" // butil::fd_guard
25#include "butil/reloadable_flags.h"25#include "butil/reloadable_flags.h"
26+#include "butil/iobuf_inl.h"
26#include "ubsocket.h"27#include "ubsocket.h"
27 28 
28namespace brpc {29namespace brpc {
@@ -50,8 +51,6 @@ DEFINE_uint32(ubiobuf_tiny_pool_threshold, DEFAULT_TINY_POOL_BLOCK_PAYLOAD_CAP,
50 "Buf size not greater than threshold will use buf of tiny pool.");51 "Buf size not greater than threshold will use buf of tiny pool.");
51BUTIL_VALIDATE_GFLAG(ubiobuf_tiny_pool_threshold, validate_ubiobuf_tiny_pool_threshold);52BUTIL_VALIDATE_GFLAG(ubiobuf_tiny_pool_threshold, validate_ubiobuf_tiny_pool_threshold);
52 53 
53-const UBIOBuf::Area UBIOBuf::INVALID_AREA;
54- 
55UBIOBuf::UBIOBuf(const Movable& rhs) {54UBIOBuf::UBIOBuf(const Movable& rhs) {
56 reset_block_ref(_sv.refs[0]);55 reset_block_ref(_sv.refs[0]);
57 reset_block_ref(_sv.refs[1]);56 reset_block_ref(_sv.refs[1]);
@@ -60,16 +59,7 @@ UBIOBuf::UBIOBuf(const Movable& rhs) {
60 59 
61namespace ubiobuf {60namespace ubiobuf {
62 61 
63-struct TLSData {62+using iobuf::TLSData;
64- // Head of the UB TLS block chain.
65- IOBuf::Block* block_head;
66- 
67- // Number of UB TLS blocks.
68- int num_blocks;
69- 
70- // True if the UB TLS block chain is registered to the thread.
71- bool registered;
72-};
73 63 
74static __thread TLSData g_ub_data = { NULL, 0, false };64static __thread TLSData g_ub_data = { NULL, 0, false };
75static __thread TLSData g_tiny_pool_data = { NULL, 0, false };65static __thread TLSData g_tiny_pool_data = { NULL, 0, false };
@@ -115,6 +105,7 @@ static butil::static_atomic<size_t> g_ub_nblock = BUTIL_STATIC_ATOMIC_INIT(0);
115static butil::static_atomic<size_t> g_ub_blockmem = BUTIL_STATIC_ATOMIC_INIT(0);105static butil::static_atomic<size_t> g_ub_blockmem = BUTIL_STATIC_ATOMIC_INIT(0);
116static butil::static_atomic<size_t> g_num_hit_ub_threshold = BUTIL_STATIC_ATOMIC_INIT(0);106static butil::static_atomic<size_t> g_num_hit_ub_threshold = BUTIL_STATIC_ATOMIC_INIT(0);
117 107 
108+IOBuf::Block* create_ub_block_with_fallback();
118void remove_tls_ub_block_chain();109void remove_tls_ub_block_chain();
119void remove_tls_tiny_pool_block_chain();110void remove_tls_tiny_pool_block_chain();
120 111 
@@ -153,6 +144,27 @@ static inline int max_blocks_per_thread() {
153 return IsIOBufProfilerEnabled() ? 0 : MAX_BLOCKS_PER_THREAD;144 return IsIOBufProfilerEnabled() ? 0 : MAX_BLOCKS_PER_THREAD;
154}145}
155 146 
147+namespace {
148+struct UBTLSBlockPolicy {
149+ static TLSData* tls_data() {
150+ return butil::ubiobuf::get_g_ub_data();
151+ }
152+ static int max_blocks_per_thread() {
153+ return butil::ubiobuf::max_blocks_per_thread();
154+ }
155+ static butil::IOBuf::Block* create_block() {
156+ return butil::ubiobuf::create_ub_block_with_fallback();
157+ }
158+ static void remove_tls_block_chain() {
159+ butil::ubiobuf::remove_tls_ub_block_chain();
160+ }
161+ static void inc_hit_tls_threshold(size_t) {
162+ butil::ubiobuf::inc_g_num_hit_ub_threshold();
163+ }
164+ static void check_block(butil::IOBuf::Block*) {}
165+};
166+} // namespace
167+ 
156size_t block_count() {168size_t block_count() {
157 return g_ub_nblock.load(butil::memory_order_relaxed);169 return g_ub_nblock.load(butil::memory_order_relaxed);
158}170}
@@ -204,26 +216,7 @@ IOBuf::Block* create_tiny_ub_block_with_fallback() {
204 return b ? b : create_ub_block_with_fallback();216 return b ? b : create_ub_block_with_fallback();
205}217}
206 218 
207-void release_tls_ub_block(IOBuf::Block* b) {219+ 
208- if (!b) {
209- return;
210- }
211- TLSData* ub_data = get_g_ub_data();
212- if (b->full()) {
213- b->dec_ref();
214- } else if (ub_data->num_blocks >= max_blocks_per_thread()) {
215- b->dec_ref();
216- inc_g_num_hit_ub_threshold();
217- } else {
218- b->u.portal_next = ub_data->block_head;
219- ub_data->block_head = b;
220- ++ub_data->num_blocks;
221- if (!ub_data->registered) {
222- ub_data->registered = true;
223- butil::thread_atexit(remove_tls_ub_block_chain);
224- }
225- }
226-}
227 220 
228void release_tls_tiny_pool_block(IOBuf::Block* b) {221void release_tls_tiny_pool_block(IOBuf::Block* b) {
229 if (!b) {222 if (!b) {
@@ -282,34 +275,7 @@ void remove_tls_tiny_pool_block_chain() {
282 ub_data.num_blocks = 0;275 ub_data.num_blocks = 0;
283}276}
284 277 
285-IOBuf::Block* share_tls_ub_block() {278+ 
286- TLSData& ub_data = g_ub_data;
287- IOBuf::Block* const b = ub_data.block_head;
288- if (b != NULL && !b->full()) {
289- return b;
290- }
291- IOBuf::Block* new_block = NULL;
292- if (b) {
293- new_block = b;
294- while (new_block && new_block->full()) {
295- IOBuf::Block* const saved_next = new_block->u.portal_next;
296- new_block->dec_ref();
297- --ub_data.num_blocks;
298- new_block = saved_next;
299- }
300- } else if (!ub_data.registered) {
301- ub_data.registered = true;
302- butil::thread_atexit(remove_tls_ub_block_chain);
303- }
304- if (!new_block) {
305- new_block = create_ub_block_with_fallback();
306- if (new_block) {
307- ++ub_data.num_blocks;
308- }
309- }
310- ub_data.block_head = new_block;
311- return new_block;
312-}
313 279 
314static inline bool is_escape_ub_block(IOBuf::Block* b) {280static inline bool is_escape_ub_block(IOBuf::Block* b) {
315 return b && (b->flags & IOBUF_BLOCK_FLAGS_UB_ESCAPE);281 return b && (b->flags & IOBUF_BLOCK_FLAGS_UB_ESCAPE);
@@ -368,39 +334,7 @@ IOBuf::Block* share_tls_tiny_pool_block() {
368 return new_block;334 return new_block;
369}335}
370 336 
371-void release_tls_ub_block_chain(IOBuf::Block* b) {337+ 
372- TLSData& ub_data = g_ub_data;
373- size_t n = 0;
374- if (ub_data.num_blocks >= max_blocks_per_thread()) {
375- do {
376- ++n;
377- IOBuf::Block* const saved_next = b->u.portal_next;
378- b->dec_ref();
379- b = saved_next;
380- } while (b);
381- inc_g_num_hit_ub_threshold();
382- return;
383- }
384- IOBuf::Block* first_b = b;
385- IOBuf::Block* last_b = NULL;
386- do {
387- ++n;
388- CHECK(!b->full());
389- CHECK(b->flags & IOBUF_BLOCK_FLAGS_UB);
390- if (b->u.portal_next == NULL) {
391- last_b = b;
392- break;
393- }
394- b = b->u.portal_next;
395- } while (true);
396- last_b->u.portal_next = ub_data.block_head;
397- ub_data.block_head = first_b;
398- ub_data.num_blocks += n;
399- if (!ub_data.registered) {
400- ub_data.registered = true;
401- butil::thread_atexit(remove_tls_ub_block_chain);
402- }
403-}
404 338 
405void release_tls_tiny_pool_block_chain(IOBuf::Block* b) {339void release_tls_tiny_pool_block_chain(IOBuf::Block* b) {
406 TLSData& ub_data = g_tiny_pool_data;340 TLSData& ub_data = g_tiny_pool_data;
@@ -437,27 +371,7 @@ void release_tls_tiny_pool_block_chain(IOBuf::Block* b) {
437 }371 }
438}372}
439 373 
440-IOBuf::Block* acquire_tls_ub_block() {374+ 
441- TLSData& ub_data = g_ub_data;
442- IOBuf::Block* b = ub_data.block_head;
443- if (!b) {
444- return create_ub_block_with_fallback();
445- }
446- while (b->full()) {
447- IOBuf::Block* const saved_next = b->u.portal_next;
448- b->dec_ref();
449- ub_data.block_head = saved_next;
450- --ub_data.num_blocks;
451- b = saved_next;
452- if (!b) {
453- return create_ub_block_with_fallback();
454- }
455- }
456- ub_data.block_head = b->u.portal_next;
457- --ub_data.num_blocks;
458- b->u.portal_next = NULL;
459- return b;
460-}
461 375 
462IOBuf::Block* acquire_tls_tiny_ub_block() {376IOBuf::Block* acquire_tls_tiny_ub_block() {
463 TLSData& ub_data = g_tiny_pool_data;377 TLSData& ub_data = g_tiny_pool_data;
@@ -495,11 +409,11 @@ static inline bool select_tiny_pool(size_t count)
495static inline IOBuf::Block* share_tls_append_block_with_fallback(size_t count)409static inline IOBuf::Block* share_tls_append_block_with_fallback(size_t count)
496{410{
497 if (!select_tiny_pool(count)) {411 if (!select_tiny_pool(count)) {
498- return share_tls_ub_block();412+ return iobuf::detail::share_tls_block<UBTLSBlockPolicy>();
499 }413 }
500 IOBuf::Block* b = share_tls_tiny_pool_block();414 IOBuf::Block* b = share_tls_tiny_pool_block();
501 // tiny block will fallback to normal or escape buf if b == NULL415 // tiny block will fallback to normal or escape buf if b == NULL
502- return b ? b : share_tls_ub_block();416+ return b ? b : iobuf::detail::share_tls_block<UBTLSBlockPolicy>();
503}417}
504 418 
505} // namespace ubiobuf419} // namespace ubiobuf
@@ -607,95 +521,6 @@ int UBIOBuf::normalize_to_tiny_pool(IOBuf* buf) {
607 return copied_blocks;521 return copied_blocks;
608}522}
609 523 
610-void UBIOBuf::operator=(const IOBuf& rhs) {
611- if (this == &rhs) {
612- return;
613- }
614- clear();
615- append(rhs);
616-}
617- 
618-void UBIOBuf::operator=(const Movable& rhs) {
619- clear();
620- append(rhs);
621-}
622- 
623-void UBIOBuf::operator=(const char* s) {
624- clear();
625- append(s);
626-}
627- 
628-void UBIOBuf::operator=(const std::string& s) {
629- clear();
630- append(s);
631-}
632- 
633-void UBIOBuf::append(const IOBuf& other) {
634- const size_t nref = other._ref_num();
635- for (size_t i = 0; i < nref; ++i) {
636- _push_back_ref(other._ref_at(i));
637- }
638-}
639- 
640-void UBIOBuf::append(const Movable& movable_other) {
641- if (empty()) {
642- swap(movable_other.value());
643- } else {
644- IOBuf& other = movable_other.value();
645- const size_t nref = other._ref_num();
646- for (size_t i = 0; i < nref; ++i) {
647- _move_back_ref(other._ref_at(i));
648- }
649- if (!other._small()) {
650- delete[] other._bv.refs;
651- }
652- reset_block_ref(other._sv.refs[0]);
653- reset_block_ref(other._sv.refs[1]);
654- }
655-}
656- 
657-int UBIOBuf::push_back(char c) {
658- Block* b = ubiobuf::share_tls_ub_block();
659- if (BAIDU_UNLIKELY(!b)) {
660- return -1;
661- }
662- b->data[b->size] = c;
663- const BlockRef r = { b->size, 1, b };
664- ++b->size;
665- _push_back_ref(r);
666- return 0;
667-}
668- 
669-int UBIOBuf::append(char const* s) {
670- if (BAIDU_LIKELY(s != NULL)) {
671- return append(s, strlen(s));
672- }
673- return -1;
674-}
675- 
676-int UBIOBuf::append(void const* data, size_t count) {
677- if (BAIDU_UNLIKELY(!data)) {
678- return -1;
679- }
680- if (count == 1) {
681- return push_back(*(char const*)data);
682- }
683- size_t total_nc = 0;
684- while (total_nc < count) {
685- Block* b = ubiobuf::share_tls_ub_block();
686- if (BAIDU_UNLIKELY(!b)) {
687- return -1;
688- }
689- const size_t nc = std::min(count - total_nc, b->left_space());
690- ubiobuf::cp(b->data + b->size, (char*)data + total_nc, nc);
691- 
692- const BlockRef r = { (uint32_t)b->size, (uint32_t)nc, b };
693- _push_back_ref(r);
694- b->size += nc;
695- total_nc += nc;
696- }
697- return 0;
698-}
699 524 
700int UBIOBuf::append_to_tiny_pool_with_fallback(void const* data, size_t count, size_t block_size) {525int UBIOBuf::append_to_tiny_pool_with_fallback(void const* data, size_t count, size_t block_size) {
701 if (BAIDU_UNLIKELY(!data)) {526 if (BAIDU_UNLIKELY(!data)) {
@@ -718,106 +543,28 @@ int UBIOBuf::append_to_tiny_pool_with_fallback(void const* data, size_t count, s
718 return 0;543 return 0;
719}544}
720 545 
721-int UBIOBuf::appendv(const const_iovec* vec, size_t n) {546+IOBuf::Block* UBIOBuf::_share_tls_block() {
722- size_t offset = 0;547+ return iobuf::detail::share_tls_block<ubiobuf::UBTLSBlockPolicy>();
723- for (size_t i = 0; i < n;) {
724- Block* b = ubiobuf::share_tls_ub_block();
725- if (BAIDU_UNLIKELY(!b)) {
726- return -1;
727- }
728- uint32_t total_cp = 0;
729- for (; i < n; ++i, offset = 0) {
730- const const_iovec& vec_i = vec[i];
731- const size_t nc = std::min(vec_i.iov_len - offset,
732- b->left_space() - total_cp);
733- ubiobuf::cp(b->data + b->size + total_cp,
734- (char*)vec_i.iov_base + offset, nc);
735- total_cp += nc;
736- offset += nc;
737- if (offset != vec_i.iov_len) {
738- break;
739- }
740- }
741- 
742- const BlockRef r = { (uint32_t)b->size, total_cp, b };
743- b->size += total_cp;
744- _push_back_ref(r);
745- }
746- return 0;
747}548}
748 549 
749-int UBIOBuf::append(const std::string& s) {550+IOBuf::Block* UBIOBuf::_create_block(uint32_t block_size) {
750- return append(s.data(), s.length());551+ const bool use_tiny_pool = ubiobuf::select_tiny_pool(block_size);
552+ return use_tiny_pool ? ubiobuf::acquire_tls_tiny_ub_block() : _acquire_tls_block();
751}553}
752 554 
753-int UBIOBuf::append_user_data_with_meta(void* data,555+IOBuf::Block* UBIOBuf::_acquire_tls_block() {
754- size_t size,556+ return iobuf::detail::acquire_tls_block<ubiobuf::UBTLSBlockPolicy>();
755- std::function<void(void*)> deleter,
756- uint64_t meta) {
757- if (size > 0xFFFFFFFFULL - 100) {
758- LOG(FATAL) << "data_size=" << size << " is too large";
759- return -1;
760- }
761- if (!deleter) {
762- deleter = ::free;
763- }
764- if (!size) {
765- deleter(data);
766- return 0;
767- }
768- char* mem = (char*)malloc(sizeof(Block) + sizeof(UserDataExtension));
769- if (mem == NULL) {
770- return -1;
771- }
772- Block* b = new (mem) Block((char*)data, size, std::move(deleter));
773- b->u.data_meta = meta;
774- const BlockRef r = { 0, b->cap, b };
775- _move_back_ref(r);
776- return 0;
777}557}
778 558 
779-int UBIOBuf::resize(size_t n, char c) {559+void UBIOBuf::_release_tls_block(IOBuf::Block* block) {
780- const size_t saved_len = length();560+ if (block == NULL) {
781- if (n < saved_len) {561+ return;
782- pop_back(saved_len - n);
783- return 0;
784 }562 }
785- const size_t count = n - saved_len;563+ if (block->flags & IOBUF_BLOCK_FLAGS_UB_TINY_POOL) {
786- size_t total_nc = 0;564+ ubiobuf::release_tls_tiny_pool_block(block);
787- while (total_nc < count) {565+ } else {
788- Block* b = ubiobuf::share_tls_ub_block();566+ iobuf::detail::release_tls_block<ubiobuf::UBTLSBlockPolicy>(block);
789- if (BAIDU_UNLIKELY(!b)) {
790- return -1;
791- }
792- const size_t nc = std::min(count - total_nc, b->left_space());
793- memset(b->data + b->size, c, nc);
794- 
795- const BlockRef r = { (uint32_t)b->size, (uint32_t)nc, b };
796- _push_back_ref(r);
797- b->size += nc;
798- total_nc += nc;
799 }567 }
800- return 0;
801-}
802- 
803-UBIOBuf::Area UBIOBuf::reserve(size_t count) {
804- Area result = INVALID_AREA;
805- size_t total_nc = 0;
806- while (total_nc < count) {
807- Block* b = ubiobuf::share_tls_ub_block();
808- if (BAIDU_UNLIKELY(!b)) {
809- return INVALID_AREA;
810- }
811- const size_t nc = std::min(count - total_nc, b->left_space());
812- const BlockRef r = { (uint32_t)b->size, (uint32_t)nc, b };
813- _push_back_ref(r);
814- if (total_nc == 0) {
815- result = make_ub_area(_ref_num() - 1, _back_ref().length - nc, count);
816- }
817- total_nc += nc;
818- b->size += nc;
819- }
820- return result;
821}568}
822 569 
823int UBIOBuf::normalize() {570int UBIOBuf::normalize() {
@@ -892,158 +639,15 @@ size_t UBIOBuf::get_block_size() {
892 return UBIOBuf::DEFAULT_BLOCK_SIZE;639 return UBIOBuf::DEFAULT_BLOCK_SIZE;
893}640}
894 641 
895-UBIOBufAsZeroCopyOutputStream::UBIOBufAsZeroCopyOutputStream(UBIOBuf* buf)
896- : _buf(buf)
897- , _block_size(0)
898- , _cur_block(NULL)
899- , _byte_count(0) {
900-}
901- 
902-UBIOBufAsZeroCopyOutputStream::UBIOBufAsZeroCopyOutputStream(
903- UBIOBuf* buf, uint32_t block_size)
904- : _buf(buf)
905- , _block_size(block_size)
906- , _cur_block(NULL)
907- , _byte_count(0) {
908-}
909- 
910-UBIOBufAsZeroCopyOutputStream::~UBIOBufAsZeroCopyOutputStream() {
911- _release_block();
912-}
913- 
914-bool UBIOBufAsZeroCopyOutputStream::Next(void** data, int* size) {
915- if (_cur_block == NULL || _cur_block->full()) {
916- _release_block();
917- if (_block_size > 0) {
918- const bool use_tiny_pool = ubiobuf::select_tiny_pool(_block_size);
919- _cur_block = use_tiny_pool ?
920- ubiobuf::acquire_tls_tiny_ub_block() :
921- ubiobuf::acquire_tls_ub_block();
922- } else {
923- _cur_block = ubiobuf::acquire_tls_ub_block();
924- }
925- if (_cur_block == NULL) {
926- return false;
927- }
928- }
929- const IOBuf::BlockRef r = { _cur_block->size,
930- (uint32_t)_cur_block->left_space(),
931- _cur_block };
932- *data = _cur_block->data + r.offset;
933- *size = r.length;
934- _cur_block->size = _cur_block->cap;
935- _buf->_push_back_ref(r);
936- _byte_count += r.length;
937- return true;
938-}
939- 
940-void UBIOBufAsZeroCopyOutputStream::BackUp(int count) {
941- while (!_buf->empty()) {
942- IOBuf::BlockRef& r = _buf->_back_ref();
943- if (_cur_block) {
944- if (r.block != _cur_block) {
945- LOG(FATAL) << "r.block=" << r.block
946- << " does not match _cur_block=" << _cur_block;
947- return;
948- }
949- if (r.offset + r.length != _cur_block->size) {
950- LOG(FATAL) << "r.offset(" << r.offset << ") + r.length("
951- << r.length << ") != _cur_block->size("
952- << _cur_block->size << ")";
953- return;
954- }
955- } else {
956- if (r.block->ref_count() == 1) {
957- if (r.offset + r.length != r.block->size) {
958- LOG(FATAL) << "r.offset(" << r.offset << ") + r.length("
959- << r.length << ") != r.block->size("
960- << r.block->size << ")";
961- return;
962- }
963- } else if (r.offset + r.length != r.block->size) {
964- _byte_count -= _buf->pop_back(count);
965- return;
966- }
967- _cur_block = r.block;
968- _cur_block->inc_ref();
969- }
970- if (BAIDU_LIKELY(r.length > (uint32_t)count)) {
971- r.length -= count;
972- if (!_buf->_small()) {
973- _buf->_bv.nbytes -= count;
974- }
975- _cur_block->size -= count;
976- _byte_count -= count;
977- _release_block();
978- return;
979- }
980- _cur_block->size -= r.length;
981- _byte_count -= r.length;
982- count -= r.length;
983- _buf->_pop_back_ref();
984- _release_block();
985- if (count == 0) {
986- return;
987- }
988- }
989- LOG_IF(FATAL, count != 0) << "BackUp an empty UBIOBuf";
990-}
991- 
992-int64_t UBIOBufAsZeroCopyOutputStream::ByteCount() const {
993- return _byte_count;
994-}
995- 
996-void UBIOBufAsZeroCopyOutputStream::_release_block() {
997- if (_cur_block == NULL) {
998- return;
999- }
1000- if (_cur_block->flags & IOBUF_BLOCK_FLAGS_UB_TINY_POOL) {
1001- ubiobuf::release_tls_tiny_pool_block(_cur_block);
1002- } else {
1003- ubiobuf::release_tls_ub_block(_cur_block);
1004- }
1005- _cur_block = NULL;
1006-}
1007- 
1008-UBIOBufAsSnappySink::UBIOBufAsSnappySink(butil::UBIOBuf& buf)
1009- : _cur_buf(NULL), _cur_len(0), _buf(&buf), _buf_stream(&buf) {
1010-}
1011- 
1012-void UBIOBufAsSnappySink::Append(const char* bytes, size_t n) {
1013- if (_cur_len > 0) {
1014- CHECK(bytes == _cur_buf && static_cast<int>(n) <= _cur_len)
1015- << "bytes must be _cur_buf";
1016- _buf_stream.BackUp(_cur_len - n);
1017- _cur_len = 0;
1018- } else {
1019- _buf->append(bytes, n);
1020- }
1021-}
1022- 
1023-char* UBIOBufAsSnappySink::GetAppendBuffer(size_t length, char* scratch) {
1024- if (length <= 8000) {
1025- if (_buf_stream.Next(reinterpret_cast<void**>(&_cur_buf), &_cur_len)) {
1026- if (_cur_len >= static_cast<int>(length)) {
1027- return _cur_buf;
1028- } else {
1029- _buf_stream.BackUp(_cur_len);
1030- }
1031- } else {
1032- LOG(FATAL) << "Fail to alloc buffer";
1033- }
1034- }
1035- _cur_buf = NULL;
1036- _cur_len = 0;
1037- return scratch;
1038-}
1039 642 
1040ssize_t IOPortal::ub_append_from_file_descriptor(643ssize_t IOPortal::ub_append_from_file_descriptor(
1041 int fd, size_t max_count) {644 int fd, size_t max_count) {
1042 if (max_count == 0) {645 if (max_count == 0) {
1043 return 0;646 return 0;
1044 }647 }
648+ 
1045 if (_block == NULL) { 649 if (_block == NULL) {
1046- _block = ubiobuf::acquire_tls_ub_block(); 650+ _block = iobuf::detail::acquire_tls_block<ubiobuf::UBTLSBlockPolicy>();
1047 if (BAIDU_UNLIKELY(!_block)) { 651 if (BAIDU_UNLIKELY(!_block)) {
1048 errno = ENOMEM; 652 errno = ENOMEM;
1049 return -1; 653 return -1;
@@ -1068,7 +672,7 @@ ssize_t IOPortal::ub_append_from_file_descriptor(
1068 } 672 }
1069#endif 673#endif
1070 if (empty()) { 674 if (empty()) {
1071- ubiobuf::release_tls_ub_block(_block);675+ iobuf::detail::release_tls_block_chain<ubiobuf::UBTLSBlockPolicy>(_block);
1072 _block = NULL;676 _block = NULL;
1073 }677 }
1074 return nr;678 return nr;
@@ -26,88 +26,32 @@ namespace butil {
26 26 
27class UBIOBuf : public IOBuf {27class UBIOBuf : public IOBuf {
28public:28public:
29- static const size_t DEFAULT_BLOCK_SIZE = IOBuf::DEFAULT_BLOCK_SIZE;
30- typedef IOBuf::Area Area;
31- static const Area INVALID_AREA = IOBuf::INVALID_AREA;
32- 
33 UBIOBuf() {}29 UBIOBuf() {}
34 UBIOBuf(const UBIOBuf& rhs) : IOBuf(rhs) {}30 UBIOBuf(const UBIOBuf& rhs) : IOBuf(rhs) {}
35 UBIOBuf(const Movable& rhs);31 UBIOBuf(const Movable& rhs);
36 ~UBIOBuf() override { clear(); }32 ~UBIOBuf() override { clear(); }
37- void operator=(const UBIOBuf& rhs) { operator=(static_cast<const IOBuf&>(rhs)); }33+ using IOBuf::operator=;
38- void operator=(const IOBuf& rhs) override;34+ 
39- void operator=(const Movable& rhs) override;
40- void operator=(const char* s) override;
41- void operator=(const std::string& s) override;
42- void append(const IOBuf& other) override;
43- void append(const Movable& other) override;
44- int push_back(char c) override;
45- int append(void const* data, size_t count) override;
46 int append_to_tiny_pool_with_fallback(void const* data, size_t count, size_t block_size);35 int append_to_tiny_pool_with_fallback(void const* data, size_t count, size_t block_size);
47- int appendv(const const_iovec vec[], size_t n) override;36+
48- int appendv(const iovec* vec, size_t n) override
49- { return appendv((const const_iovec*)vec, n); }
50- int append(char const* s) override;
51- int append(const std::string& s) override;
52- int append_user_data(void* data, size_t size, std::function<void(void*)> deleter) override {
53- return append_user_data_with_meta(data, size, std::move(deleter), 0);
54- }
55- int append_user_data_with_meta(void* data, size_t size,
56- std::function<void(void*)> deleter,
57- uint64_t meta) override;
58- int resize(size_t n) override { return resize(n, '\0'); }
59- int resize(size_t n, char c) override;
60- Area reserve(size_t n) override;
61 bool use_ub() const override { return true; }37 bool use_ub() const override { return true; }
62 int normalize();38 int normalize();
63 static int normalize(IOBuf* buf);39 static int normalize(IOBuf* buf);
64 static bool has_ub_block(const IOBuf* buf);40 static bool has_ub_block(const IOBuf* buf);
65 static size_t get_block_size();41 static size_t get_block_size();
66 42 
43+protected:
44+ Block* _share_tls_block() override;
45+ Block* _create_block(uint32_t block_size) override;
46+ Block* _acquire_tls_block() override;
47+ void _release_tls_block(Block* block) override;
48+ 
67private:49private:
68 static int append_to_tiny_pool(UBIOBuf* out, const void* data, size_t len);50 static int append_to_tiny_pool(UBIOBuf* out, const void* data, size_t len);
69 static int append_to_registered_ub_pool(UBIOBuf* out, const void* data, size_t len, bool use_tiny_pool);51 static int append_to_registered_ub_pool(UBIOBuf* out, const void* data, size_t len, bool use_tiny_pool);
70 static int normalize_to_tiny_pool(IOBuf* buf);52 static int normalize_to_tiny_pool(IOBuf* buf);
71};53};
72 54 
73-class UBIOBufAsZeroCopyOutputStream
74- : public google::protobuf::io::ZeroCopyOutputStream {
75-public:
76- explicit UBIOBufAsZeroCopyOutputStream(UBIOBuf*);
77- UBIOBufAsZeroCopyOutputStream(UBIOBuf*, uint32_t block_size);
78- ~UBIOBufAsZeroCopyOutputStream();
79- 
80- bool Next(void** data, int* size) override;
81- void BackUp(int count) override;
82- int64_t ByteCount() const override;
83- 
84-private:
85- void _release_block();
86- 
87- UBIOBuf* _buf;
88- uint32_t _block_size;
89- IOBuf::Block *_cur_block;
90- int64_t _byte_count;
91-};
92- 
93-// Wrap UBIOBuf into output of snappy compression.
94-class UBIOBufAsSnappySink : public butil::snappy::Sink {
95-public:
96- explicit UBIOBufAsSnappySink(butil::UBIOBuf& buf);
97- virtual ~UBIOBufAsSnappySink() {}
98- 
99- // Append "bytes[0,n-1]" to this.
100- void Append(const char* bytes, size_t n) override;
101- 
102- // Returns a writable buffer of the specified length for appending.
103- char* GetAppendBuffer(size_t length, char* scratch) override;
104- 
105-private:
106- char* _cur_buf;
107- int _cur_len;
108- butil::UBIOBuf* _buf;
109- butil::UBIOBufAsZeroCopyOutputStream _buf_stream;
110-};
111 55 
112namespace ubiobuf {56namespace ubiobuf {
113 57