1 #pragma once
2
3 #include <string>
4 #include <unordered_map>
5 #include <mutex>
6 #include <random>
7 #include <ctime>
8
9 #ifdef FI_FLAG
10
11 // Utility class for failure injection
12 class FIUtil {
13 public:
14 // Failure points table
15 class FPTable {
16 public:
17 void Update(const std::string& name, int percent) {
18 mu_.lock();
19 failure_points_[name] = percent;
20 mu_.unlock();
21 }
22
23 void Remove(const std::string& name) {
24 mu_.lock();
25 failure_points_.erase(name);
26 mu_.unlock();
27 }
28
29 int LookUp(const std::string& name) {
30 std::lock_guard<std::mutex> guard(mu_);
31 auto iter = failure_points_.find(name);
32 if (iter == failure_points_.end()) {
33 return -1; // Not found
34 }
35 return iter->second;
36 }
37
38 void Clear() {
39 std::unordered_map<std::string, int> tmp;
40 mu_.lock();
41 failure_points_.swap(tmp);
42 mu_.unlock();
43 }
44
45 private:
46 std::mutex mu_;
47 std::unordered_map<std::string, int> failure_points_;
48 };
49
50 static inline FPTable* GetFPTable() {
51 static FPTable table;
52 return &table;
53 }
54
55 static inline int Random() {
56 static std::default_random_engine g(std::time(NULL));
57 std::uniform_int_distribution<int> d(0,99);
58 return d(g);
59 }
60 };
61
62 #define FI_RETURN(name, retval) \
63 do { \
64 if (FIUtil::Random() < FIUtil::GetFPTable()->LookUp(name)) { \
65 return retval; \
66 } \
67 } while (0)
68
69 #define FI_ASSIGN(name, var, val) \
70 do { \
71 if (FIUtil::Random() < FIUtil::GetFPTable()->LookUp(name)) { \
72 var = val; \
73 } \
74 } while (0)
75
76 #define FI_ENABLE(name) \
77 do { \
78 FIUtil::GetFPTable()->Update(name, 100); \
79 } while (0)
80
81 #define FI_ENABLE_RANDOM(name, percent) \
82 do { \
83 FIUtil::GetFPTable()->Update(name, percent); \
84 } while (0)
85
86 #define FI_DISABLE(name) \
87 do { \
88 FIUtil::GetFPTable()->Update(name, 0); \
89 } while (0)
90
91 #define FI_DISABLE_ALL() \
92 do { \
93 FIUtil::GetFPTable()->Clear(); \
94 } while (0)
95
96 #else
97
98 #define FI_RETURN(name, retval)
99 #define FI_ASSIGN(name, var, val)
100 #define FI_ENABLE(name)
101 #define FI_ENABLE_RANDOM(name, percent)
102 #define FI_DISABLE(name)
103 #define FI_DISABLE_ALL()
104
105 #endif // FI_FLAG
