package main import ( "encoding/binary" "log" "strconv" "time" "github.com/boltdb/bolt" "github.com/hamcha/clessy/tg" ) const ( MessageTypeText int = 0 MessageTypeAudio int = 1 MessageTypePhoto int = 2 MessageTypeSticker int = 3 MessageTypeVideo int = 4 MessageTypeVoice int = 5 MessageTypeContact int = 6 MessageTypeLocation int = 7 MessageTypeDocument int = 8 MessageTypeMax int = 9 ) type Stats struct { ByUserCount map[string]uint64 ByWeekday [7]uint64 ByHour [24]uint64 ByType [MessageTypeMax]uint64 TodayDate time.Time Today uint64 TotalCount uint64 Replies uint64 Forward uint64 } var stats Stats func MakeUint(bval []byte, bucketName string, key string) uint64 { if bval != nil { intval, bts := binary.Uvarint(bval) if bts > 0 { return intval } else { log.Printf("[%s] Value of key \"%s\" is NaN: %v\r\n", bucketName, key, bval) return 0 } } else { log.Printf("[%s] Key \"%s\" does not exist, set to 0\n", bucketName, key) return 0 } } func PutUint(value uint64) []byte { bytes := make([]byte, 10) n := binary.PutUvarint(bytes, value) return bytes[:n] } func loadStats() { // Load today stats.TodayDate = time.Now() err := db.Update(func(tx *bolt.Tx) error { b, err := tx.CreateBucketIfNotExists([]byte("global")) if err != nil { return err } // Load total messages counter stats.TotalCount = MakeUint(b.Get([]byte("count")), "global", "count") // Load total replies counter stats.Replies = MakeUint(b.Get([]byte("replies")), "global", "replies") // Load total replies counter stats.Forward = MakeUint(b.Get([]byte("forward")), "global", "forward") // Load hour counters b, err = tx.CreateBucketIfNotExists([]byte("hour")) if err != nil { return err } for i := 0; i < 24; i++ { stats.ByHour[i] = MakeUint(b.Get([]byte{byte(i)}), "hour", strconv.Itoa(i)) } // Load weekday counters b, err = tx.CreateBucketIfNotExists([]byte("weekday")) if err != nil { return err } for i := 0; i < 7; i++ { stats.ByWeekday[i] = MakeUint(b.Get([]byte{byte(i)}), "weekday", strconv.Itoa(i)) } // Load today's message counter, if possible b, err = tx.CreateBucketIfNotExists([]byte("date")) if err != nil { return err } todayKey := stats.TodayDate.Format("2006-1-2") stats.Today = MakeUint(b.Get([]byte(todayKey)), "date", todayKey) // Load user counters stats.ByUserCount = make(map[string]uint64) b, err = tx.CreateBucketIfNotExists([]byte("users-count")) if err != nil { return err } b.ForEach(func(user, messages []byte) error { stats.ByUserCount[string(user)] = MakeUint(messages, "users-count", string(user)) return nil }) // Load type counters b, err = tx.CreateBucketIfNotExists([]byte("types")) if err != nil { return err } for i := 0; i < MessageTypeMax; i++ { stats.ByType[i] = MakeUint(b.Get([]byte{byte(i)}), "types", strconv.Itoa(i)) } return nil }) assert(err) } func updateDate() { err := db.Update(func(tx *bolt.Tx) error { b := tx.Bucket([]byte("date")) err := b.Put([]byte(stats.TodayDate.Format("2006-1-2")), PutUint(stats.Today)) if err != nil { return err } return nil }) if err != nil { log.Println("[updateDate] Couldn't save last day stats: " + err.Error()) } stats.TodayDate = time.Now() stats.Today = 0 } func updateMean(currentMean, meanCount, newValue uint64) uint64 { return ((currentMean * meanCount) + newValue) / (meanCount + 1) } func updateStats(message tg.APIMessage) { // // Local update // // DB Update flags updatetype := 0 updatereplies := false updateforward := false // Update total count stats.TotalCount++ // Update individual user's count username := message.User.Username val, exists := stats.ByUserCount[username] if !exists { val = 0 } stats.ByUserCount[username] = val + 1 // Update time counters now := time.Now() hour := now.Hour() wday := now.Weekday() stats.ByHour[hour]++ stats.ByWeekday[wday]++ // Check for day reset if now.Day() != stats.TodayDate.Day() { updateDate() } stats.Today++ // Text message if message.Text != nil { stats.ByType[MessageTypeText]++ updatetype = MessageTypeText } // Audio message if message.Audio != nil { stats.ByType[MessageTypeAudio]++ updatetype = MessageTypeAudio } // Photo if message.Photo != nil { stats.ByType[MessageTypePhoto]++ updatetype = MessageTypePhoto } // Sticker if message.Sticker != nil { stats.ByType[MessageTypeSticker]++ updatetype = MessageTypeSticker } // Video if message.Video != nil { stats.ByType[MessageTypeVideo]++ updatetype = MessageTypeVideo } // Voice message if message.Voice != nil { stats.ByType[MessageTypeVoice]++ updatetype = MessageTypeVoice } // Contact if message.Contact != nil { stats.ByType[MessageTypeContact]++ updatetype = MessageTypeContact } // Location if message.Location != nil { stats.ByType[MessageTypeLocation]++ updatetype = MessageTypeLocation } // Document if message.Document != nil { stats.ByType[MessageTypeDocument]++ updatetype = MessageTypeDocument } // Reply if message.ReplyTo != nil { stats.Replies++ updatereplies = true } // Forwarded message if message.FwdUser != nil { stats.Forward++ updateforward = true } // // DB Update // err := db.Update(func(tx *bolt.Tx) error { // Update total counters b := tx.Bucket([]byte("global")) err := b.Put([]byte("count"), PutUint(stats.TotalCount)) if err != nil { return err } if updatereplies { err = b.Put([]byte("replies"), PutUint(stats.Replies)) if err != nil { return err } } if updateforward { err = b.Put([]byte("forward"), PutUint(stats.Forward)) if err != nil { return err } } // Update time counters b = tx.Bucket([]byte("hour")) err = b.Put([]byte{byte(hour)}, PutUint(stats.ByHour[hour])) if err != nil { return err } b = tx.Bucket([]byte("weekday")) err = b.Put([]byte{byte(wday)}, PutUint(stats.ByWeekday[wday])) if err != nil { return err } b = tx.Bucket([]byte("date")) todayKey := stats.TodayDate.Format("2006-1-2") err = b.Put([]byte(todayKey), PutUint(stats.Today)) if err != nil { return err } // Update user counters b = tx.Bucket([]byte("users-count")) err = b.Put([]byte(username), PutUint(stats.ByUserCount[username])) if err != nil { return err } // Update type counter b = tx.Bucket([]byte("types")) err = b.Put([]byte{byte(updatetype)}, PutUint(stats.ByType[updatetype])) if err != nil { return err } return nil }) if err != nil { log.Println("[updateStats] Got error while updating DB: " + err.Error()) } }