You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 

161 lines
3.6 KiB

  1. package main
  2. import (
  3. "context"
  4. "fmt"
  5. "log"
  6. "time"
  7. "github.com/go-redis/redis/v8"
  8. )
  9. type ProjectBackfeedManager struct {
  10. Context context.Context
  11. Cancel context.CancelFunc
  12. Done chan bool
  13. C chan *BackfeedItem
  14. Name string
  15. BackfeedRedis *redis.ClusterClient
  16. ProjectRedis *redis.Client
  17. //Lock sync.RWMutex
  18. ProjectConfig ProjectConfig
  19. }
  20. type ProjectRedisConfig struct {
  21. Host string `json:"host"`
  22. Pass string `json:"pass"`
  23. Port int `json:"port"`
  24. }
  25. func (that *ProjectBackfeedManager) RedisConfigDiffers(new *ProjectRedisConfig) bool {
  26. if that.ProjectConfig.RedisConfig == nil && new == nil {
  27. return false
  28. }
  29. return that.ProjectConfig.RedisConfig == nil || new == nil || *that.ProjectConfig.RedisConfig != *new
  30. }
  31. func (that *ProjectBackfeedManager) PushItem(ctx context.Context, item *BackfeedItem) error {
  32. //that.Lock.RLock()
  33. //defer that.Lock.RUnlock()
  34. //if that.C == nil {
  35. // return false
  36. //}
  37. select {
  38. case <-ctx.Done():
  39. return ctx.Err()
  40. case <-that.Context.Done():
  41. return fmt.Errorf("backfeed channel closed")
  42. case that.C <- item:
  43. return nil
  44. //default:
  45. // return fmt.Errorf("backfeed channel full")
  46. }
  47. }
  48. func (that *ProjectBackfeedManager) PopItem(blocking bool) (*BackfeedItem, bool) {
  49. if blocking {
  50. select {
  51. case <-that.Context.Done():
  52. return nil, false
  53. case item, ok := <-that.C:
  54. return item, ok
  55. }
  56. } else {
  57. select {
  58. case <-that.Context.Done():
  59. return nil, false
  60. case item, ok := <-that.C:
  61. return item, ok
  62. default:
  63. return nil, false
  64. }
  65. }
  66. }
  67. func (that *ProjectBackfeedManager) Do() {
  68. defer close(that.Done)
  69. //defer that.CloseItemChannel()
  70. defer that.Cancel()
  71. pipe := that.BackfeedRedis.Pipeline()
  72. for {
  73. select {
  74. case <-that.Context.Done():
  75. break
  76. case <-that.Done:
  77. break
  78. default:
  79. }
  80. item, ok := that.PopItem(true)
  81. if !ok {
  82. break
  83. }
  84. keyMap := map[string][][]byte{}
  85. key := fmt.Sprintf("%s:%02x:%s", that.Name, item.PrimaryShard, item.SecondaryShard)
  86. keyMap[key] = append(keyMap[key], item.Item)
  87. wrapped := 1
  88. for wrapped < ItemWrapSize {
  89. item, ok := that.PopItem(false)
  90. if !ok {
  91. break
  92. }
  93. key := fmt.Sprintf("%s:%02x:%s", that.Name, item.PrimaryShard, item.SecondaryShard)
  94. keyMap[key] = append(keyMap[key], item.Item)
  95. wrapped++
  96. }
  97. select {
  98. case <-that.Context.Done():
  99. break
  100. case <-that.Done:
  101. break
  102. default:
  103. }
  104. now := time.Now()
  105. resultMap := map[string]*redis.Cmd{}
  106. lastTS := make([]any, 0, len(keyMap)*2)
  107. for key := range keyMap {
  108. lastTS = append(lastTS, key)
  109. lastTS = append(lastTS, fmt.Sprintf("%d", now.Unix()))
  110. }
  111. pipe.HSet(context.Background(), ":last_ts", lastTS...)
  112. for key, items := range keyMap {
  113. args := []any{
  114. "bf.madd",
  115. key,
  116. }
  117. for _, item := range items {
  118. args = append(args, item)
  119. }
  120. resultMap[key] = pipe.Do(context.Background(), args...)
  121. }
  122. if _, err := pipe.Exec(context.Background()); err != nil {
  123. log.Printf("%s", err)
  124. }
  125. var sAddItems []any
  126. for key, items := range keyMap {
  127. res, err := resultMap[key].BoolSlice()
  128. if err != nil {
  129. log.Printf("%s", err)
  130. continue
  131. }
  132. if len(res) != len(keyMap[key]) {
  133. continue
  134. }
  135. for i, v := range res {
  136. if v {
  137. sAddItems = append(sAddItems, items[i])
  138. }
  139. }
  140. }
  141. dupes := wrapped - len(sAddItems)
  142. if len(sAddItems) != 0 {
  143. if err := that.ProjectRedis.SAdd(context.Background(), fmt.Sprintf("%s:todo:backfeed", that.Name), sAddItems...).Err(); err != nil {
  144. log.Printf("failed to sadd items for %s: %s", that.Name, err)
  145. }
  146. }
  147. if dupes > 0 {
  148. that.BackfeedRedis.HIncrBy(context.Background(), ":", that.Name, int64(dupes))
  149. }
  150. }
  151. }