thrift/processor/StatsProcessor.h

0001 /*
0002  * Licensed to the Apache Software Foundation (ASF) under one
0003  * or more contributor license agreements. See the NOTICE file
0004  * distributed with this work for additional information
0005  * regarding copyright ownership. The ASF licenses this file
0006  * to you under the Apache License, Version 2.0 (the
0007  * "License"); you may not use this file except in compliance
0008  * with the License. You may obtain a copy of the License at
0009  *
0010  *   http://www.apache.org/licenses/LICENSE-2.0
0011  *
0012  * Unless required by applicable law or agreed to in writing,
0013  * software distributed under the License is distributed on an
0014  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
0015  * KIND, either express or implied. See the License for the
0016  * specific language governing permissions and limitations
0017  * under the License.
0018  */
0019
0020 #ifndef STATSPROCESSOR_H
0021 #define STATSPROCESSOR_H
0022
0023 #include <memory>
0024 #include <thrift/transport/TTransport.h>
0025 #include <thrift/protocol/TProtocol.h>
0026 #include <TProcessor.h>
0027
0028 namespace apache {
0029 namespace thrift {
0030 namespace processor {
0031
0032 /*
0033  * Class for keeping track of function call statistics and printing them if desired
0034  *
0035  */
0036 class StatsProcessor : public apache::thrift::TProcessor {
0037 public:
0038   StatsProcessor(bool print, bool frequency) : print_(print), frequency_(frequency) {}
0039   virtual ~StatsProcessor(){};
0040
0041   virtual bool process(std::shared_ptr<apache::thrift::protocol::TProtocol> piprot,
0042                        std::shared_ptr<apache::thrift::protocol::TProtocol> poprot,
0043                        void* serverContext) {
0044
0045     piprot_ = piprot;
0046
0047     std::string fname;
0048     apache::thrift::protocol::TMessageType mtype;
0049     int32_t seqid;
0050
0051     piprot_->readMessageBegin(fname, mtype, seqid);
0052     if (mtype != apache::thrift::protocol::T_CALL && mtype != apache::thrift::protocol::T_ONEWAY) {
0053       if (print_) {
0054         printf("Unknown message type\n");
0055       }
0056       throw apache::thrift::TException("Unexpected message type");
0057     }
0058     if (print_) {
0059       printf("%s (", fname.c_str());
0060     }
0061     if (frequency_) {
0062       if (frequency_map_.find(fname) != frequency_map_.end()) {
0063         frequency_map_[fname]++;
0064       } else {
0065         frequency_map_[fname] = 1;
0066       }
0067     }
0068
0069     apache::thrift::protocol::TType ftype;
0070     int16_t fid;
0071
0072     while (true) {
0073       piprot_->readFieldBegin(fname, ftype, fid);
0074       if (ftype == apache::thrift::protocol::T_STOP) {
0075         break;
0076       }
0077
0078       printAndPassToBuffer(ftype);
0079       if (print_) {
0080         printf(", ");
0081       }
0082     }
0083
0084     if (print_) {
0085       printf("\b\b)\n");
0086     }
0087     return true;
0088   }
0089
0090   const std::map<std::string, int64_t>& get_frequency_map() { return frequency_map_; }
0091
0092 protected:
0093   void printAndPassToBuffer(apache::thrift::protocol::TType ftype) {
0094     switch (ftype) {
0095     case apache::thrift::protocol::T_BOOL: {
0096       bool boolv;
0097       piprot_->readBool(boolv);
0098       if (print_) {
0099         printf("%d", boolv);
0100       }
0101     } break;
0102     case apache::thrift::protocol::T_BYTE: {
0103       int8_t bytev;
0104       piprot_->readByte(bytev);
0105       if (print_) {
0106         printf("%d", bytev);
0107       }
0108     } break;
0109     case apache::thrift::protocol::T_I16: {
0110       int16_t i16;
0111       piprot_->readI16(i16);
0112       if (print_) {
0113         printf("%d", i16);
0114       }
0115     } break;
0116     case apache::thrift::protocol::T_I32: {
0117       int32_t i32;
0118       piprot_->readI32(i32);
0119       if (print_) {
0120         printf("%d", i32);
0121       }
0122     } break;
0123     case apache::thrift::protocol::T_I64: {
0124       int64_t i64;
0125       piprot_->readI64(i64);
0126       if (print_) {
0127         printf("%ld", i64);
0128       }
0129     } break;
0130     case apache::thrift::protocol::T_DOUBLE: {
0131       double dub;
0132       piprot_->readDouble(dub);
0133       if (print_) {
0134         printf("%f", dub);
0135       }
0136     } break;
0137     case apache::thrift::protocol::T_STRING: {
0138       std::string str;
0139       piprot_->readString(str);
0140       if (print_) {
0141         printf("%s", str.c_str());
0142       }
0143     } break;
0144     case apache::thrift::protocol::T_STRUCT: {
0145       std::string name;
0146       int16_t fid;
0147       apache::thrift::protocol::TType ftype;
0148       piprot_->readStructBegin(name);
0149       if (print_) {
0150         printf("<");
0151       }
0152       while (true) {
0153         piprot_->readFieldBegin(name, ftype, fid);
0154         if (ftype == apache::thrift::protocol::T_STOP) {
0155           break;
0156         }
0157         printAndPassToBuffer(ftype);
0158         if (print_) {
0159           printf(",");
0160         }
0161         piprot_->readFieldEnd();
0162       }
0163       piprot_->readStructEnd();
0164       if (print_) {
0165         printf("\b>");
0166       }
0167     } break;
0168     case apache::thrift::protocol::T_MAP: {
0169       apache::thrift::protocol::TType keyType;
0170       apache::thrift::protocol::TType valType;
0171       uint32_t i, size;
0172       piprot_->readMapBegin(keyType, valType, size);
0173       if (print_) {
0174         printf("{");
0175       }
0176       for (i = 0; i < size; i++) {
0177         printAndPassToBuffer(keyType);
0178         if (print_) {
0179           printf("=>");
0180         }
0181         printAndPassToBuffer(valType);
0182         if (print_) {
0183           printf(",");
0184         }
0185       }
0186       piprot_->readMapEnd();
0187       if (print_) {
0188         printf("\b}");
0189       }
0190     } break;
0191     case apache::thrift::protocol::T_SET: {
0192       apache::thrift::protocol::TType elemType;
0193       uint32_t i, size;
0194       piprot_->readSetBegin(elemType, size);
0195       if (print_) {
0196         printf("{");
0197       }
0198       for (i = 0; i < size; i++) {
0199         printAndPassToBuffer(elemType);
0200         if (print_) {
0201           printf(",");
0202         }
0203       }
0204       piprot_->readSetEnd();
0205       if (print_) {
0206         printf("\b}");
0207       }
0208     } break;
0209     case apache::thrift::protocol::T_LIST: {
0210       apache::thrift::protocol::TType elemType;
0211       uint32_t i, size;
0212       piprot_->readListBegin(elemType, size);
0213       if (print_) {
0214         printf("[");
0215       }
0216       for (i = 0; i < size; i++) {
0217         printAndPassToBuffer(elemType);
0218         if (print_) {
0219           printf(",");
0220         }
0221       }
0222       piprot_->readListEnd();
0223       if (print_) {
0224         printf("\b]");
0225       }
0226     } break;
0227     default:
0228       break;
0229     }
0230   }
0231
0232   std::shared_ptr<apache::thrift::protocol::TProtocol> piprot_;
0233   std::map<std::string, int64_t> frequency_map_;
0234
0235   bool print_;
0236   bool frequency_;
0237 };
0238 }
0239 }
0240 } // apache::thrift::processor
0241
0242 #endif