support checksum

Author: yanglimingcn

example/echo_c++/client.cpp

@@ -31,6 +31,7 @@ DEFINE_string(load_balancer, "", "The algorithm for load balancing");
 DEFINE_int32(timeout_ms, 100, "RPC timeout in milliseconds");
 DEFINE_int32(max_retry, 3, "Max retries(not including the first RPC)"); 
 DEFINE_int32(interval_ms, 1000, "Milliseconds between consecutive requests");
+DEFINE_bool(enable_checksum, false, "Enable checksum or not");
 
 int main(int argc, char* argv[]) {
     // Parse gflags. We recommend you to use gflags as well.
@@ -71,6 +72,11 @@ int main(int argc, char* argv[]) {
         // being serialized into protobuf messages.
         cntl.request_attachment().append(FLAGS_attachment);
 
+        // Use checksum, only support CRC32C now.
+        if (FLAGS_enable_checksum) {
+            cntl.set_request_checksum_type(brpc::CHECKSUM_TYPE_CRC32C);
+        }
+
         // Because `done'(last parameter) is NULL, this function waits until
         // the response comes back or error occurs(including timedout).
         stub.Echo(&cntl, &request, &response, NULL);

example/echo_c++/server.cpp

@@ -29,6 +29,7 @@ DEFINE_string(listen_addr, "", "Server listen address, may be IPV4/IPV6/UDS."
             " If this is set, the flag port will be ignored");
 DEFINE_int32(idle_timeout_s, -1, "Connection will be closed if there is no "
              "read/write operations during the last `idle_timeout_s'");
+DEFINE_bool(enable_checksum, false, "Enable checksum or not");
 
 // Your implementation of example::EchoService
 // Notice that implementing brpc::Describable grants the ability to put
@@ -75,6 +76,11 @@ class EchoServiceImpl : public EchoService {
             // being serialized into protobuf messages.
             cntl->response_attachment().append(cntl->request_attachment());
         }
+
+        // Use checksum, only support CRC32C now.
+        if (FLAGS_enable_checksum) {
+            cntl->set_response_checksum_type(brpc::CHECKSUM_TYPE_CRC32C);
+        }
     }
 
     // optional

src/brpc/checksum.cpp

@@ -0,0 +1,100 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "brpc/checksum.h"
+
+#include "brpc/protocol.h"
+#include "butil/logging.h"
+
+namespace brpc {
+
+static const int MAX_HANDLER_SIZE = 1024;
+static ChecksumHandler s_handler_map[MAX_HANDLER_SIZE] = {{NULL, NULL, NULL}};
+
+int RegisterChecksumHandler(ChecksumType type, ChecksumHandler handler) {
+    if (NULL == handler.Compute) {
+        LOG(FATAL) << "Invalid parameter: handler function is NULL";
+        return -1;
+    }
+    int index = type;
+    if (index < 0 || index >= MAX_HANDLER_SIZE) {
+        LOG(FATAL) << "ChecksumType=" << type << " is out of range";
+        return -1;
+    }
+    if (s_handler_map[index].Compute != NULL) {
+        LOG(FATAL) << "ChecksumType=" << type << " was registered";
+        return -1;
+    }
+    s_handler_map[index] = handler;
+    return 0;
+}
+
+// Find ChecksumHandler by type.
+// Returns NULL if not found
+inline const ChecksumHandler* FindChecksumHandler(ChecksumType type) {
+    int index = type;
+    if (index < 0 || index >= MAX_HANDLER_SIZE) {
+        LOG(ERROR) << "ChecksumType=" << type << " is out of range";
+        return NULL;
+    }
+    if (NULL == s_handler_map[index].Compute) {
+        return NULL;
+    }
+    return &s_handler_map[index];
+}
+
+const char* ChecksumTypeToCStr(ChecksumType type) {
+    if (type == CHECKSUM_TYPE_NONE) {
+        return "none";
+    }
+    const ChecksumHandler* handler = FindChecksumHandler(type);
+    return (handler != NULL ? handler->name : "unknown");
+}
+
+void ListChecksumHandler(std::vector<ChecksumHandler>* vec) {
+    vec->clear();
+    for (int i = 0; i < MAX_HANDLER_SIZE; ++i) {
+        if (s_handler_map[i].Compute != NULL) {
+            vec->push_back(s_handler_map[i]);
+        }
+    }
+}
+
+// Compute `data' checksum
+void ComputeDataChecksum(const ChecksumIn& in, ChecksumType checksum_type) {
+    if (checksum_type == CHECKSUM_TYPE_NONE) {
+        return;
+    }
+    const ChecksumHandler* handler = FindChecksumHandler(checksum_type);
+    if (NULL != handler) {
+        handler->Compute(in);
+    }
+}
+
+// Verify `data' checksum Returns true on success, false otherwise
+bool VerifyDataChecksum(const ChecksumIn& in, ChecksumType checksum_type) {
+    if (checksum_type == CHECKSUM_TYPE_NONE) {
+        return true;
+    }
+    const ChecksumHandler* handler = FindChecksumHandler(checksum_type);
+    if (NULL != handler) {
+        return handler->Verify(in);
+    }
+    return true;
+}
+
+}  // namespace brpc

src/brpc/checksum.h

@@ -0,0 +1,65 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#ifndef BRPC_CHECKSUM_H
+#define BRPC_CHECKSUM_H
+
+#include <google/protobuf/message.h>  // Message
+
+#include "brpc/controller.h"
+#include "brpc/options.pb.h"  // ChecksumType
+#include "butil/iobuf.h"      // butil::IOBuf
+
+namespace brpc {
+
+struct ChecksumIn {
+    const butil::IOBuf* buf;
+    Controller* cntl;
+};
+
+struct ChecksumHandler {
+    // checksum `buf'.
+    // Returns checksum value
+    void (*Compute)(const ChecksumIn& in);
+
+    // verify buf checksum
+    // Rerturn true on success, false otherwise
+    bool (*Verify)(const ChecksumIn& in);
+
+    // Name of the checksum algorithm, must be string constant.
+    const char* name;
+};
+
+// [NOT thread-safe] Register `handler' using key=`type'
+// Returns 0 on success, -1 otherwise
+int RegisterChecksumHandler(ChecksumType type, ChecksumHandler handler);
+
+// Returns the `name' of the checksumType if registered
+const char* ChecksumTypeToCStr(ChecksumType type);
+
+// Put all registered handlers into `vec'.
+void ListChecksumHandler(std::vector<ChecksumHandler>* vec);
+
+// Compute `data' checksum and set to controller
+void ComputeDataChecksum(const ChecksumIn& in, ChecksumType checksum_type);
+
+// Verify `data' checksum Returns true on success, false otherwise
+bool VerifyDataChecksum(const ChecksumIn& in, ChecksumType checksum_type);
+
+}  // namespace brpc
+
+#endif  // BRPC_CHECKSUM_H

src/brpc/controller.cpp

@@ -269,6 +269,8 @@ void Controller::ResetPods() {
     _preferred_index = -1;
     _request_compress_type = COMPRESS_TYPE_NONE;
     _response_compress_type = COMPRESS_TYPE_NONE;
+    _request_checksum_type = CHECKSUM_TYPE_NONE;
+    _response_checksum_type = CHECKSUM_TYPE_NONE;
     _fail_limit = UNSET_MAGIC_NUM;
     _pipelined_count = 0;
     _inheritable.Reset();
@@ -1346,6 +1348,7 @@ void Controller::SaveClientSettings(ClientSettings* s) const {
     s->tos = _tos;
     s->connection_type = _connection_type;
     s->request_compress_type = _request_compress_type;
+    s->request_checksum_type = _request_checksum_type;
     s->log_id = log_id();
     s->has_request_code = has_request_code();
     s->request_code = _request_code;
@@ -1359,6 +1362,7 @@ void Controller::ApplyClientSettings(const ClientSettings& s) {
     set_type_of_service(s.tos);
     set_connection_type(s.connection_type);
     set_request_compress_type(s.request_compress_type);
+    set_request_checksum_type(s.request_checksum_type);
     set_log_id(s.log_id);
     set_flag(FLAGS_REQUEST_CODE, s.has_request_code);
     _request_code = s.request_code;

src/brpc/controller.h

@@ -241,6 +241,9 @@ friend void policy::ProcessThriftRequest(InputMessageBase*);
     // Set compression method for request.
     void set_request_compress_type(CompressType t) { _request_compress_type = t; }
 
+    // Set checksum type for request.
+    void set_request_checksum_type(ChecksumType t) { _request_checksum_type = t; }
+
     // Required by some load balancers.
     void set_request_code(uint64_t request_code) {
         add_flag(FLAGS_REQUEST_CODE);
@@ -464,6 +467,9 @@ friend void policy::ProcessThriftRequest(InputMessageBase*);
 
     // Set compression method for response.
     void set_response_compress_type(CompressType t) { _response_compress_type = t; }
+
+    // Set checksum type for response.
+    void set_response_checksum_type(ChecksumType t) { _response_checksum_type = t; }
 
     // Non-zero when this RPC call is traced (by rpcz or rig).
     // NOTE: Only valid at server-side, always zero at client-side.
@@ -552,6 +558,8 @@ friend void policy::ProcessThriftRequest(InputMessageBase*);
     const std::string& request_id() const { return _inheritable.request_id; }
     CompressType request_compress_type() const { return _request_compress_type; }
     CompressType response_compress_type() const { return _response_compress_type; }
+    ChecksumType request_checksum_type() const { return _request_checksum_type; }
+    ChecksumType response_checksum_type() const { return _response_checksum_type; }
     const HttpHeader& http_request() const 
     { return _http_request != NULL ? *_http_request : DefaultHttpHeader(); }
 
@@ -693,6 +701,7 @@ friend void policy::ProcessThriftRequest(InputMessageBase*);
         int32_t tos;
         ConnectionType connection_type;         
         CompressType request_compress_type;
+        ChecksumType request_checksum_type;
         uint64_t log_id;
         bool has_request_code;
         int64_t request_code;
@@ -834,6 +843,9 @@ friend void policy::ProcessThriftRequest(InputMessageBase*);
     int _preferred_index;
     CompressType _request_compress_type;
     CompressType _response_compress_type;
+    ChecksumType _request_checksum_type;
+    ChecksumType _response_checksum_type;
+    std::string _checksum_value;
     Inheritable _inheritable;
     int _pchan_sub_count;
     google::protobuf::Message* _response;

src/brpc/details/controller_private_accessor.h

@@ -152,6 +152,16 @@ class ControllerPrivateAccessor {
         return *this;
     }
 
+    void set_checksum_value(const char* c, size_t size) {
+        _cntl->_checksum_value.assign(c, size);
+    }
+
+    void set_checksum_value(const std::string& c) {
+        _cntl->_checksum_value = c;
+    }
+
+    const std::string& checksum_value() const { return _cntl->_checksum_value; }
+
 private:
     Controller* _cntl;
 };

src/brpc/global.cpp

@@ -60,6 +60,10 @@
 #include "brpc/policy/gzip_compress.h"
 #include "brpc/policy/snappy_compress.h"
 
+// Checksum handlers
+#include "brpc/checksum.h"
+#include "brpc/policy/crc32c_checksum.h"
+
 // Protocols
 #include "brpc/protocol.h"
 #include "brpc/policy/baidu_rpc_protocol.h"
@@ -401,6 +405,13 @@ static void GlobalInitializeOrDieImpl() {
         exit(1);
     }
 
+    // Checksum Handlers
+    const ChecksumHandler crc32c_checksum = {Crc32cCompute, Crc32cVerify,
+                                             "crc32c"};
+    if (RegisterChecksumHandler(CHECKSUM_TYPE_CRC32C, crc32c_checksum) != 0) {
+        exit(1);
+    }
+
     // Protocols
     Protocol baidu_protocol = { ParseRpcMessage,
                                 SerializeRpcRequest, PackRpcRequest,

src/brpc/options.proto

@@ -74,6 +74,11 @@ enum CompressType {
     COMPRESS_TYPE_LZ4 = 4;
 }
 
+enum ChecksumType {
+    CHECKSUM_TYPE_NONE = 0;
+    CHECKSUM_TYPE_CRC32C = 1;
+}
+
 enum ContentType {
     CONTENT_TYPE_PB = 0;
     CONTENT_TYPE_JSON = 1;

src/brpc/policy/baidu_rpc_meta.proto

@@ -34,6 +34,8 @@ message RpcMeta {
     optional StreamSettings stream_settings = 8;
     map<string, string> user_fields = 9;
     optional ContentType content_type = 10;
+    optional int32 checksum_type = 11;
+    optional bytes checksum_value = 12;
 }
 
 message RpcRequestMeta {