package main import ( "bufio" "compress/flate" "compress/gzip" "context" "encoding/json" "fmt" "io" "log" "net/http" "strings" "sync" "time" "github.com/go-redis/redis/v8" "github.com/gorilla/mux" "github.com/tevino/abool/v2" ) type GlobalBackfeedManager struct { Context context.Context Cancel context.CancelFunc ActiveFeeds map[string]*ProjectBackfeedManager ActiveSlugs map[string]string TrackerRedis *redis.Client BackfeedRedis *redis.ClusterClient Lock sync.RWMutex Populated *abool.AtomicBool } func (that *GlobalBackfeedManager) RefreshFeeds() error { slugProjectMap, err := that.TrackerRedis.HGetAll(that.Context, "backfeed").Result() if err != nil { return err } var projects []string projectSlugMap := map[string][]string{} for slug, project := range slugProjectMap { projectSlugMap[project] = append(projectSlugMap[project], slug) } for project := range projectSlugMap { projects = append(projects, project) } projectConfigs := map[string]ProjectConfig{} if len(projects) != 0 { cfgi, err := that.TrackerRedis.HMGet(that.Context, "trackers", projects...).Result() if err != nil { return err } if len(projects) != len(cfgi) { return fmt.Errorf("hmget result had unexpected length") } for i, project := range projects { configString, ok := cfgi[i].(string) if !ok { continue } config := ProjectConfig{} if err := json.Unmarshal([]byte(configString), &config); err != nil { continue } projectConfigs[project] = config } } projects = nil for project := range projectSlugMap { if _, has := projectConfigs[project]; !has { delete(projectSlugMap, project) continue } projects = append(projects, project) } for slug, project := range slugProjectMap { if _, has := projectConfigs[project]; !has { delete(slugProjectMap, slug) } } // add feeds for new projects for _, project := range projects { projectConfig := projectConfigs[project] var outdatedProjectBackfeedManager *ProjectBackfeedManager if projectBackfeedManager, has := that.ActiveFeeds[project]; has { if that.ActiveFeeds[project].RedisConfigDiffers(projectConfig.RedisConfig) { outdatedProjectBackfeedManager = projectBackfeedManager } else { continue } } ctx, cancel := context.WithCancel(that.Context) projectBackfeedManager := &ProjectBackfeedManager{ Context: ctx, Cancel: cancel, Done: make(chan bool), C: make(chan *BackfeedItem, ItemChannelBuffer), BackfeedRedis: that.BackfeedRedis, Name: project, ProjectConfig: projectConfig, } if projectConfig.RedisConfig != nil { projectBackfeedManager.ProjectRedis = redis.NewClient(&redis.Options{ Addr: fmt.Sprintf("%s:%d", projectConfig.RedisConfig.Host, projectConfig.RedisConfig.Port), Username: "default", Password: projectConfig.RedisConfig.Pass, ReadTimeout: 15 * time.Minute, }) } else { projectBackfeedManager.ProjectRedis = that.TrackerRedis } go projectBackfeedManager.Do() that.Lock.Lock() that.ActiveFeeds[project] = projectBackfeedManager that.Lock.Unlock() if outdatedProjectBackfeedManager != nil { outdatedProjectBackfeedManager.Cancel() <-outdatedProjectBackfeedManager.Done log.Printf("updated project: %s", project) } else { log.Printf("added project: %s", project) } } that.Lock.Lock() that.ActiveSlugs = slugProjectMap that.Lock.Unlock() // remove feeds for old projects for project, projectBackfeedManager := range that.ActiveFeeds { if _, has := projectSlugMap[project]; has { continue } log.Printf("removing project: %s", project) that.Lock.Lock() delete(that.ActiveFeeds, project) that.Lock.Unlock() projectBackfeedManager.Cancel() <-projectBackfeedManager.Done log.Printf("removed project: %s", project) } if !that.Populated.IsSet() { that.Populated.Set() } return nil } func (that *GlobalBackfeedManager) GetFeed(slug string) *ProjectBackfeedManager { that.Lock.RLock() defer that.Lock.RUnlock() project, has := that.ActiveSlugs[slug] if !has { return nil } projectBackfeedManager, has := that.ActiveFeeds[project] if !has { return nil } return projectBackfeedManager } func (that *GlobalBackfeedManager) HandleLegacy(res http.ResponseWriter, req *http.Request) { defer req.Body.Close() vars := mux.Vars(req) slug := vars["slug"] secondaryShard := req.URL.Query().Get("shard") if strings.ContainsAny(secondaryShard, ":/") { WriteResponse(res, http.StatusBadRequest, fmt.Errorf("invalid shard name")) return } projectBackfeedManager := that.GetFeed(slug) if projectBackfeedManager == nil { WriteResponse(res, http.StatusNotFound, fmt.Errorf("%s", "no such backfeed channel")) return } splitter := &Splitter{ Delimiter: []byte(req.URL.Query().Get("delimiter")), IgnoreEOF: req.URL.Query().Get("ignoreeof") != "", } if len(splitter.Delimiter) == 0 { splitter.Delimiter = []byte{0x00} } var body io.ReadCloser switch req.Header.Get("Content-Encoding") { case "": body = req.Body case "gzip": var err error body, err = gzip.NewReader(req.Body) if err != nil { WriteResponse(res, http.StatusBadRequest, err) return } defer body.Close() case "deflate": body = flate.NewReader(req.Body) defer body.Close() default: WriteResponse(res, http.StatusBadRequest, fmt.Errorf("unsupported Content-Encoding: %s", req.Header.Get("Content-Encoding"))) } scanner := bufio.NewScanner(body) scanner.Split(splitter.Split) n := 0 for scanner.Scan() { b := scanner.Bytes() if len(b) == 0 { continue } bcopy := make([]byte, len(b)) copy(bcopy, b) item := &BackfeedItem{ PrimaryShard: GenShardHash(bcopy), SecondaryShard: secondaryShard, Item: bcopy, } if err := projectBackfeedManager.PushItem(req.Context(), item); err != nil { WriteResponse(res, http.StatusInternalServerError, err) return } n++ } if err := scanner.Err(); err != nil { WriteResponse(res, http.StatusBadRequest, err) return } WriteResponse(res, http.StatusOK, fmt.Sprintf("%d items queued for deduplication", n)) return } func (that *GlobalBackfeedManager) HandleHealth(res http.ResponseWriter, req *http.Request) { if that.Populated.IsNotSet() { WriteResponse(res, http.StatusServiceUnavailable, fmt.Errorf("%s", "backfeed not populated")) return } if err := that.BackfeedRedis.ForEachShard(req.Context(), func(ctx context.Context, client *redis.Client) error { client.ClientGetName(ctx) return client.Ping(ctx).Err() }); err != nil { WriteResponse(res, http.StatusInternalServerError, fmt.Errorf("failed to ping backfeed redis: %s", err)) return } WriteResponse(res, http.StatusOK, "ok") } func (that *GlobalBackfeedManager) HandlePing(res http.ResponseWriter, _ *http.Request) { WriteResponse(res, http.StatusOK, "pong") } func (that *GlobalBackfeedManager) HandleUnlink(res http.ResponseWriter, req *http.Request) { vars := mux.Vars(req) key := vars["key"] if strings.Count(key, ":") < 2 { WriteResponse(res, http.StatusBadRequest, fmt.Errorf("invalid key format")) return } lock := sync.Mutex{} keys := []string{} if err := that.BackfeedRedis.ForEachShard(req.Context(), func(ctx context.Context, client *redis.Client) error { cursor := uint64(0) var shardKeys []string for { var err error var keysBatch []string keysBatch, cursor, err = client.Scan(ctx, cursor, key, 1000).Result() if err != nil && err != redis.Nil { return err } shardKeys = append(shardKeys, keysBatch...) if cursor == 0 { break } } lock.Lock() defer lock.Unlock() keys = append(keys, shardKeys...) return nil }); err != nil && err != redis.Nil { WriteResponse(res, http.StatusInternalServerError, err) return } pipe := that.BackfeedRedis.Pipeline() pipe.HDel(req.Context(), ":last_ts", keys...) for _, key := range keys { pipe.Unlink(req.Context(), key) } if _, err := pipe.Exec(req.Context()); err != nil && err != redis.Nil { WriteResponse(res, http.StatusInternalServerError, err) return } WriteResponse(res, http.StatusOK, keys) } func (that *GlobalBackfeedManager) HandleRedisInfo(res http.ResponseWriter, req *http.Request) { infos := map[string]string{} lock := sync.Mutex{} if err := that.BackfeedRedis.ForEachShard(req.Context(), func(ctx context.Context, client *redis.Client) error { if info, err := client.Info(ctx, "all").Result(); err != nil && err != redis.Nil { return err } else { lock.Lock() defer lock.Unlock() infos[client.String()] = info } return nil }); err != nil { WriteResponse(res, http.StatusInternalServerError, err) return } WriteResponse(res, http.StatusOK, infos) } func (that *GlobalBackfeedManager) CancelAllFeeds() { that.Populated.UnSet() that.Cancel() for project, projectBackfeedManager := range that.ActiveFeeds { log.Printf("waiting for %s channel to shut down...", project) <-projectBackfeedManager.Done delete(that.ActiveFeeds, project) } }