util.go 6.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362
  1. package common
  2. import (
  3. "bytes"
  4. "encoding/json"
  5. "fmt"
  6. "github.com/go-yaml/yaml"
  7. "github.com/patrickmn/go-cache"
  8. "github.com/sirupsen/logrus"
  9. "io/ioutil"
  10. "os"
  11. "path"
  12. "path/filepath"
  13. "time"
  14. )
  15. const (
  16. logFileName = "stream.log"
  17. etc_dir = "/etc/"
  18. data_dir = "/data/"
  19. log_dir = "/log/"
  20. )
  21. var (
  22. Log *logrus.Logger
  23. Config *XStreamConf
  24. IsTesting bool
  25. logFile *os.File
  26. mockTicker *MockTicker
  27. mockTimer *MockTimer
  28. mockNow int64
  29. )
  30. type logRedirect struct {
  31. }
  32. func (l *logRedirect) Errorf(f string, v ...interface{}) {
  33. Log.Error(fmt.Sprintf(f, v...))
  34. }
  35. func (l *logRedirect) Infof(f string, v ...interface{}) {
  36. Log.Info(fmt.Sprintf(f, v...))
  37. }
  38. func (l *logRedirect) Warningf(f string, v ...interface{}) {
  39. Log.Warning(fmt.Sprintf(f, v...))
  40. }
  41. func (l *logRedirect) Debugf(f string, v ...interface{}) {
  42. Log.Debug(fmt.Sprintf(f, v...))
  43. }
  44. func LoadConf(confName string) ([]byte, error) {
  45. confDir, err := GetConfLoc()
  46. if err != nil {
  47. return nil, err
  48. }
  49. file := confDir + confName
  50. b, err := ioutil.ReadFile(file)
  51. if err != nil {
  52. return nil, err
  53. }
  54. return b, nil
  55. }
  56. type XStreamConf struct {
  57. Debug bool `yaml:"debug"`
  58. Port int `yaml:"port"`
  59. }
  60. var StreamConf = "kuiper.yaml"
  61. func init(){
  62. Log = logrus.New()
  63. Log.SetFormatter(&logrus.TextFormatter{
  64. DisableColors: true,
  65. FullTimestamp: true,
  66. })
  67. b, err := LoadConf(StreamConf)
  68. if err != nil {
  69. Log.Fatal(err)
  70. }
  71. var cfg map[string]XStreamConf
  72. if err := yaml.Unmarshal(b, &cfg); err != nil {
  73. Log.Fatal(err)
  74. }
  75. if c, ok := cfg["basic"]; !ok{
  76. Log.Fatal("no basic config in kuiper.yaml")
  77. }else{
  78. Config = &c
  79. }
  80. if !Config.Debug {
  81. logDir, err := GetLoc(log_dir)
  82. if err != nil {
  83. Log.Fatal(err)
  84. }
  85. file := logDir + logFileName
  86. logFile, err := os.OpenFile(file, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
  87. if err == nil {
  88. Log.Out = logFile
  89. } else {
  90. Log.Infof("Failed to log to file, using default stderr")
  91. }
  92. }else{
  93. Log.SetLevel(logrus.DebugLevel)
  94. }
  95. }
  96. type KeyValue interface {
  97. Open() error
  98. Close() error
  99. Set(key string, value interface{}) error
  100. Get(key string) (interface{}, bool)
  101. Delete(key string) error
  102. Keys() (keys []string, err error)
  103. }
  104. type SimpleKVStore struct {
  105. path string
  106. c *cache.Cache;
  107. }
  108. var stores = make(map[string]*SimpleKVStore)
  109. func GetSimpleKVStore(path string) *SimpleKVStore {
  110. if s, ok := stores[path]; ok {
  111. return s
  112. } else {
  113. c := cache.New(cache.NoExpiration, 0)
  114. if _, err := os.Stat(path); os.IsNotExist(err) {
  115. os.MkdirAll(path, os.ModePerm)
  116. }
  117. sStore := &SimpleKVStore{path: path + "/stores.data", c: c}
  118. stores[path] = sStore
  119. return sStore
  120. }
  121. }
  122. func (m *SimpleKVStore) Open() error {
  123. if _, err := os.Stat(m.path); os.IsNotExist(err) {
  124. return nil
  125. }
  126. if e := m.c.LoadFile(m.path); e != nil {
  127. return e
  128. }
  129. return nil
  130. }
  131. func (m *SimpleKVStore) Close() error {
  132. e := m.saveToFile()
  133. m.c.Flush() //Delete all of the values from memory.
  134. return e
  135. }
  136. func (m *SimpleKVStore) saveToFile() error {
  137. if e := m.c.SaveFile(m.path); e != nil {
  138. return e
  139. }
  140. return nil
  141. }
  142. func (m *SimpleKVStore) Set(key string, value interface{}) error {
  143. if m.c == nil {
  144. return fmt.Errorf("Cache %s has not been initialized yet.", m.path)
  145. }
  146. m.c.Set(key, value, cache.NoExpiration)
  147. return m.saveToFile()
  148. }
  149. func (m *SimpleKVStore) Get(key string) (interface{}, bool) {
  150. return m.c.Get(key)
  151. }
  152. func (m *SimpleKVStore) Delete(key string) error {
  153. m.c.Delete(key)
  154. return m.saveToFile()
  155. }
  156. func (m *SimpleKVStore) Keys() (keys []string, err error) {
  157. if m.c == nil {
  158. return nil, fmt.Errorf("Cache %s has not been initialized yet.", m.path)
  159. }
  160. its := m.c.Items()
  161. keys = make([]string, 0, len(its))
  162. for k := range its {
  163. keys = append(keys, k)
  164. }
  165. return keys, nil
  166. }
  167. func PrintMap(m map[string]string, buff *bytes.Buffer) {
  168. for k, v := range m {
  169. buff.WriteString(fmt.Sprintf("%s: %s\n", k, v))
  170. }
  171. }
  172. func CloseLogger(){
  173. if logFile != nil {
  174. logFile.Close()
  175. }
  176. }
  177. func GetConfLoc()(string, error){
  178. return GetLoc(etc_dir)
  179. }
  180. func GetDataLoc() (string, error) {
  181. return GetLoc(data_dir)
  182. }
  183. func GetLoc(subdir string)(string, error) {
  184. dir, err := os.Getwd()
  185. if err != nil {
  186. return "", err
  187. }
  188. confDir := dir + subdir
  189. if _, err := os.Stat(confDir); os.IsNotExist(err) {
  190. lastdir := dir
  191. for len(dir) > 0 {
  192. dir = filepath.Dir(dir)
  193. if lastdir == dir {
  194. break
  195. }
  196. confDir = dir + subdir
  197. if _, err := os.Stat(confDir); os.IsNotExist(err) {
  198. lastdir = dir
  199. continue
  200. } else {
  201. //Log.Printf("Trying to load file from %s", confDir)
  202. return confDir, nil
  203. }
  204. }
  205. } else {
  206. //Log.Printf("Trying to load file from %s", confDir)
  207. return confDir, nil
  208. }
  209. return "", fmt.Errorf("conf dir not found")
  210. }
  211. func GetAndCreateDataLoc(dir string) (string, error) {
  212. dataDir, err := GetDataLoc()
  213. if err != nil {
  214. return "", err
  215. }
  216. d := path.Join(path.Dir(dataDir), dir)
  217. if _, err := os.Stat(d); os.IsNotExist(err) {
  218. err = os.MkdirAll(d, 0755)
  219. if err != nil {
  220. return "", err
  221. }
  222. }
  223. return d, nil
  224. }
  225. //Time related. For Mock
  226. func GetTicker(duration int) Ticker {
  227. if IsTesting{
  228. if mockTicker == nil{
  229. mockTicker = NewMockTicker(duration)
  230. }else{
  231. mockTicker.SetDuration(duration)
  232. }
  233. return mockTicker
  234. }else{
  235. return NewDefaultTicker(duration)
  236. }
  237. }
  238. func GetTimer(duration int) Timer {
  239. if IsTesting{
  240. if mockTimer == nil{
  241. mockTimer = NewMockTimer(duration)
  242. }else{
  243. mockTimer.SetDuration(duration)
  244. }
  245. return mockTimer
  246. }else{
  247. return NewDefaultTimer(duration)
  248. }
  249. }
  250. func GetNowInMilli() int64{
  251. if IsTesting {
  252. return GetMockNow()
  253. }else{
  254. return TimeToUnixMilli(time.Now())
  255. }
  256. }
  257. func ProcessPath(p string) (string, error) {
  258. if abs, err := filepath.Abs(p); err != nil {
  259. return "", nil
  260. } else {
  261. if _, err := os.Stat(abs); os.IsNotExist(err) {
  262. return "", err;
  263. }
  264. return abs, nil
  265. }
  266. }
  267. /****** For Test Only ********/
  268. func GetMockTicker() *MockTicker{
  269. return mockTicker
  270. }
  271. func ResetMockTicker(){
  272. if mockTicker != nil{
  273. mockTicker.lastTick = 0
  274. }
  275. }
  276. func GetMockTimer() *MockTimer{
  277. return mockTimer
  278. }
  279. func SetMockNow(now int64){
  280. mockNow = now
  281. }
  282. func GetMockNow() int64{
  283. return mockNow
  284. }
  285. /*********** Type Cast Utilities *****/
  286. //TODO datetime type
  287. func ToString(input interface{}) string{
  288. return fmt.Sprintf("%v", input)
  289. }
  290. func ToInt(input interface{}) (int, error){
  291. switch t := input.(type) {
  292. case float64:
  293. return int(t), nil
  294. case int64:
  295. return int(t), nil
  296. case int:
  297. return t, nil
  298. default:
  299. return 0, fmt.Errorf("unsupported type %T of %[1]v", input)
  300. }
  301. }
  302. /*
  303. * Convert a map into a struct. The output parameter must be a pointer to a struct
  304. * The struct can have the json meta data
  305. */
  306. func MapToStruct(input map[string]interface{}, output interface{}) error{
  307. // convert map to json
  308. jsonString, err := json.Marshal(input)
  309. if err != nil{
  310. return err
  311. }
  312. // convert json to struct
  313. return json.Unmarshal(jsonString, output)
  314. }