No puede seleccionar más de 25 temas Los temas deben comenzar con una letra o número, pueden incluir guiones ('-') y pueden tener hasta 35 caracteres de largo.
 
 

210 líneas
5.2 KiB

  1. package main
  2. import (
  3. "context"
  4. "fmt"
  5. "log"
  6. "time"
  7. "github.com/go-redis/redis/v8"
  8. "github.com/hashicorp/go-multierror"
  9. )
  10. type ProjectBackfeedManager struct {
  11. Context context.Context
  12. Cancel context.CancelFunc
  13. Done chan bool
  14. C chan *BackfeedItem
  15. Name string
  16. BackfeedRedis *redis.ClusterClient
  17. ProjectRedis *redis.Client
  18. //Lock sync.RWMutex
  19. ProjectConfig ProjectConfig
  20. }
  21. type ProjectRedisConfig struct {
  22. Host string `json:"host"`
  23. Pass string `json:"pass"`
  24. Port int `json:"port"`
  25. }
  26. func (that *ProjectBackfeedManager) RedisConfigDiffers(new *ProjectRedisConfig) bool {
  27. if that.ProjectConfig.RedisConfig == nil && new == nil {
  28. return false
  29. }
  30. return that.ProjectConfig.RedisConfig == nil || new == nil || *that.ProjectConfig.RedisConfig != *new
  31. }
  32. func (that *ProjectBackfeedManager) PushItem(ctx context.Context, item *BackfeedItem) error {
  33. //that.Lock.RLock()
  34. //defer that.Lock.RUnlock()
  35. //if that.C == nil {
  36. // return false
  37. //}
  38. select {
  39. case <-ctx.Done():
  40. return ctx.Err()
  41. case <-that.Context.Done():
  42. return fmt.Errorf("backfeed channel closed")
  43. case that.C <- item:
  44. return nil
  45. //default:
  46. // return fmt.Errorf("backfeed channel full")
  47. }
  48. }
  49. func (that *ProjectBackfeedManager) PopItem(blocking bool) (*BackfeedItem, bool) {
  50. if blocking {
  51. select {
  52. case <-that.Context.Done():
  53. return nil, false
  54. case item, ok := <-that.C:
  55. return item, ok
  56. }
  57. } else {
  58. select {
  59. case <-that.Context.Done():
  60. return nil, false
  61. case item, ok := <-that.C:
  62. return item, ok
  63. default:
  64. return nil, false
  65. }
  66. }
  67. }
  68. var Tag = struct{}{}
  69. func (that *ProjectBackfeedManager) Do() {
  70. defer close(that.Done)
  71. defer that.Cancel()
  72. pipe := that.BackfeedRedis.Pipeline()
  73. for {
  74. select {
  75. case <-that.Context.Done():
  76. break
  77. case <-that.Done:
  78. break
  79. default:
  80. }
  81. queueKeyMap := map[string]map[string][][]byte{}
  82. sAddQueueItems := map[string][]any{}
  83. skipFeedQueueItems := map[string]map[string]struct{}{}
  84. wrapped := 0
  85. for wrapped < ItemWrapSize {
  86. item, ok := that.PopItem(wrapped == 0)
  87. if !ok {
  88. break
  89. }
  90. if item.SkipBloom {
  91. sAddQueueItems[item.Queue] = append(sAddQueueItems[item.Queue], item.Item)
  92. } else {
  93. key := fmt.Sprintf("%s:%02x:%s", that.Name, item.PrimaryShard, item.SecondaryShard)
  94. queueKeyMap[item.Queue][key] = append(queueKeyMap[item.Queue][key], item.Item)
  95. if item.SkipFeed {
  96. skipFeedQueueItems[item.Queue][string(item.Item)] = Tag
  97. }
  98. }
  99. wrapped++
  100. }
  101. if wrapped == 0 {
  102. break
  103. }
  104. if len(queueKeyMap) > 0 {
  105. for queue, keyMap := range queueKeyMap {
  106. lastTS := make([]any, 0, len(keyMap)*2)
  107. now := time.Now()
  108. for key := range keyMap {
  109. lastTS = append(lastTS, key)
  110. lastTS = append(lastTS, fmt.Sprintf("%d", now.Unix()))
  111. }
  112. resultMap := map[string]*redis.Cmd{}
  113. pipe.HSet(context.Background(), ":last_ts", lastTS...)
  114. try := 0
  115. for {
  116. for key, items := range keyMap {
  117. args := []any{
  118. "bf.madd",
  119. key,
  120. }
  121. for _, item := range items {
  122. args = append(args, item)
  123. }
  124. resultMap[key] = pipe.Do(context.Background(), args...)
  125. }
  126. if _, err := pipe.Exec(context.Background()); err != nil {
  127. log.Printf("%s", err)
  128. }
  129. var err error
  130. for key, items := range keyMap {
  131. res, cmdErr := resultMap[key].BoolSlice()
  132. if cmdErr != nil {
  133. err = multierror.Append(err, cmdErr)
  134. continue
  135. }
  136. if len(res) != len(keyMap[key]) {
  137. err = multierror.Append(err, fmt.Errorf("invalid response length for %s: %d != %d", key, len(res), len(keyMap[key])))
  138. continue
  139. }
  140. for i, v := range res {
  141. if v {
  142. sAddQueueItems[queue] = append(sAddQueueItems[queue], items[i])
  143. }
  144. }
  145. delete(keyMap, key)
  146. }
  147. if err == nil {
  148. break
  149. }
  150. log.Printf("%s", err)
  151. time.Sleep(time.Duration(try) * time.Second)
  152. try++
  153. }
  154. delete(queueKeyMap, queue)
  155. }
  156. }
  157. items_len := 0
  158. for _, sAddItems := range(sAddQueueItems) {
  159. items_len = items_len + len(sAddItems)
  160. }
  161. dupes := wrapped - items_len
  162. if len(sAddQueueItems) > 0 && len(skipFeedQueueItems) > 0 {
  163. for queue, sAddItems := range sAddQueueItems {
  164. skipFeedItems, exists := skipFeedQueueItems[queue]
  165. if !exists {
  166. continue
  167. }
  168. itemFiltered := make([]any, 0, len(sAddItems))
  169. for _, item := range sAddItems {
  170. itemBytes := item.([]byte)
  171. itemString := string(itemBytes)
  172. if _, exists := skipFeedItems[itemString]; !exists {
  173. itemFiltered = append(itemFiltered, item)
  174. }
  175. }
  176. sAddQueueItems[queue] = itemFiltered
  177. delete(skipFeedQueueItems, queue)
  178. }
  179. }
  180. if len(sAddQueueItems) > 0 {
  181. for queue, sAddItems := range sAddQueueItems {
  182. try := 0
  183. for {
  184. if err := that.ProjectRedis.SAdd(context.Background(), fmt.Sprintf("%s:%s", that.Name, queue), sAddItems...).Err(); err != nil {
  185. log.Printf("failed to sadd items for %s: %s", that.Name, err)
  186. time.Sleep(time.Duration(try) * time.Second)
  187. try++
  188. } else {
  189. break
  190. }
  191. }
  192. delete(sAddQueueItems, queue)
  193. }
  194. }
  195. if dupes > 0 {
  196. that.BackfeedRedis.HIncrBy(context.Background(), ":", that.Name, int64(dupes))
  197. }
  198. }
  199. }