util.go 6.2 KB

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