util.go 6.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371
  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. if err := m.c.Add(key, value, cache.NoExpiration); err != nil {
  147. return err
  148. }
  149. return m.saveToFile()
  150. }
  151. func (m *SimpleKVStore) Get(key string) (interface{}, bool) {
  152. return m.c.Get(key)
  153. }
  154. func (m *SimpleKVStore) Delete(key string) error {
  155. if m.c == nil {
  156. return fmt.Errorf("cache %s has not been initialized yet", m.path)
  157. }
  158. if _, found := m.c.Get(key); found {
  159. m.c.Delete(key)
  160. }else{
  161. return fmt.Errorf("%s is not found", key)
  162. }
  163. return m.saveToFile()
  164. }
  165. func (m *SimpleKVStore) Keys() (keys []string, err error) {
  166. if m.c == nil {
  167. return nil, fmt.Errorf("Cache %s has not been initialized yet.", m.path)
  168. }
  169. its := m.c.Items()
  170. keys = make([]string, 0, len(its))
  171. for k := range its {
  172. keys = append(keys, k)
  173. }
  174. return keys, nil
  175. }
  176. func PrintMap(m map[string]string, buff *bytes.Buffer) {
  177. for k, v := range m {
  178. buff.WriteString(fmt.Sprintf("%s: %s\n", k, v))
  179. }
  180. }
  181. func CloseLogger(){
  182. if logFile != nil {
  183. logFile.Close()
  184. }
  185. }
  186. func GetConfLoc()(string, error){
  187. return GetLoc(etc_dir)
  188. }
  189. func GetDataLoc() (string, error) {
  190. return GetLoc(data_dir)
  191. }
  192. func GetLoc(subdir string)(string, error) {
  193. dir, err := os.Getwd()
  194. if err != nil {
  195. return "", err
  196. }
  197. confDir := dir + subdir
  198. if _, err := os.Stat(confDir); os.IsNotExist(err) {
  199. lastdir := dir
  200. for len(dir) > 0 {
  201. dir = filepath.Dir(dir)
  202. if lastdir == dir {
  203. break
  204. }
  205. confDir = dir + subdir
  206. if _, err := os.Stat(confDir); os.IsNotExist(err) {
  207. lastdir = dir
  208. continue
  209. } else {
  210. //Log.Printf("Trying to load file from %s", confDir)
  211. return confDir, nil
  212. }
  213. }
  214. } else {
  215. //Log.Printf("Trying to load file from %s", confDir)
  216. return confDir, nil
  217. }
  218. return "", fmt.Errorf("conf dir not found")
  219. }
  220. func GetAndCreateDataLoc(dir string) (string, error) {
  221. dataDir, err := GetDataLoc()
  222. if err != nil {
  223. return "", err
  224. }
  225. d := path.Join(path.Dir(dataDir), dir)
  226. if _, err := os.Stat(d); os.IsNotExist(err) {
  227. err = os.MkdirAll(d, 0755)
  228. if err != nil {
  229. return "", err
  230. }
  231. }
  232. return d, nil
  233. }
  234. //Time related. For Mock
  235. func GetTicker(duration int) Ticker {
  236. if IsTesting{
  237. if mockTicker == nil{
  238. mockTicker = NewMockTicker(duration)
  239. }else{
  240. mockTicker.SetDuration(duration)
  241. }
  242. return mockTicker
  243. }else{
  244. return NewDefaultTicker(duration)
  245. }
  246. }
  247. func GetTimer(duration int) Timer {
  248. if IsTesting{
  249. if mockTimer == nil{
  250. mockTimer = NewMockTimer(duration)
  251. }else{
  252. mockTimer.SetDuration(duration)
  253. }
  254. return mockTimer
  255. }else{
  256. return NewDefaultTimer(duration)
  257. }
  258. }
  259. func GetNowInMilli() int64{
  260. if IsTesting {
  261. return GetMockNow()
  262. }else{
  263. return TimeToUnixMilli(time.Now())
  264. }
  265. }
  266. func ProcessPath(p string) (string, error) {
  267. if abs, err := filepath.Abs(p); err != nil {
  268. return "", nil
  269. } else {
  270. if _, err := os.Stat(abs); os.IsNotExist(err) {
  271. return "", err;
  272. }
  273. return abs, nil
  274. }
  275. }
  276. /****** For Test Only ********/
  277. func GetMockTicker() *MockTicker{
  278. return mockTicker
  279. }
  280. func ResetMockTicker(){
  281. if mockTicker != nil{
  282. mockTicker.lastTick = 0
  283. }
  284. }
  285. func GetMockTimer() *MockTimer{
  286. return mockTimer
  287. }
  288. func SetMockNow(now int64){
  289. mockNow = now
  290. }
  291. func GetMockNow() int64{
  292. return mockNow
  293. }
  294. /*********** Type Cast Utilities *****/
  295. //TODO datetime type
  296. func ToString(input interface{}) string{
  297. return fmt.Sprintf("%v", input)
  298. }
  299. func ToInt(input interface{}) (int, error){
  300. switch t := input.(type) {
  301. case float64:
  302. return int(t), nil
  303. case int64:
  304. return int(t), nil
  305. case int:
  306. return t, nil
  307. default:
  308. return 0, fmt.Errorf("unsupported type %T of %[1]v", input)
  309. }
  310. }
  311. /*
  312. * Convert a map into a struct. The output parameter must be a pointer to a struct
  313. * The struct can have the json meta data
  314. */
  315. func MapToStruct(input map[string]interface{}, output interface{}) error{
  316. // convert map to json
  317. jsonString, err := json.Marshal(input)
  318. if err != nil{
  319. return err
  320. }
  321. // convert json to struct
  322. return json.Unmarshal(jsonString, output)
  323. }