package tgc import ( "context" "strings" "sync" "time" "github.com/divyam234/teldrive/internal/config" "github.com/divyam234/teldrive/internal/kv" "github.com/gotd/td/telegram" "go.uber.org/zap" ) type BotWorker struct { mu sync.RWMutex bots map[int64][]string currIdx map[int64]int } func NewBotWorker() *BotWorker { return &BotWorker{ bots: make(map[int64][]string), currIdx: make(map[int64]int), } } func (w *BotWorker) Set(bots []string, channelID int64) { w.mu.Lock() defer w.mu.Unlock() w.bots[channelID] = bots w.currIdx[channelID] = 0 } func (w *BotWorker) Next(channelID int64) (string, int) { w.mu.RLock() defer w.mu.RUnlock() bots := w.bots[channelID] index := w.currIdx[channelID] w.currIdx[channelID] = (index + 1) % len(bots) return bots[index], index } type ClientStatus int const ( StatusIdle ClientStatus = iota StatusBusy ) type Client struct { Tg *telegram.Client Stop StopFunc Status ClientStatus UserID string } type StreamWorker struct { mu sync.RWMutex clients map[string]*Client currIdx map[int64]int channelBots map[int64][]string cnf *config.TGConfig kv kv.KV ctx context.Context logger *zap.SugaredLogger activeStreams int cancel context.CancelFunc } func NewStreamWorker(cnf *config.Config, kv kv.KV, logger *zap.SugaredLogger) *StreamWorker { ctx, cancel := context.WithCancel(context.Background()) worker := &StreamWorker{ cnf: &cnf.TG, kv: kv, ctx: ctx, clients: make(map[string]*Client), currIdx: make(map[int64]int), channelBots: make(map[int64][]string), logger: logger, cancel: cancel, } go worker.startIdleClientMonitor() return worker } func (w *StreamWorker) Set(bots []string, channelID int64) { w.mu.Lock() defer w.mu.Unlock() w.channelBots[channelID] = bots w.currIdx[channelID] = 0 } func (w *StreamWorker) Next(channelID int64) (*Client, error) { w.mu.Lock() defer w.mu.Unlock() bots := w.channelBots[channelID] index := w.currIdx[channelID] token := bots[index] userID := strings.Split(token, ":")[0] client, err := w.getOrCreateClient(userID, token) if err != nil { return nil, err } w.currIdx[channelID] = (index + 1) % len(bots) return client, nil } func (w *StreamWorker) IncActiveStream() error { w.mu.Lock() defer w.mu.Unlock() w.activeStreams++ return nil } func (w *StreamWorker) DecActiveStreams() error { w.mu.Lock() defer w.mu.Unlock() if w.activeStreams == 0 { return nil } w.activeStreams-- return nil } func (w *StreamWorker) getOrCreateClient(userID, token string) (*Client, error) { client, ok := w.clients[userID] if !ok || (client.Status == StatusIdle && client.Stop == nil) { middlewares := Middlewares(w.cnf, 5) tgClient, _ := BotClient(w.ctx, w.kv, w.cnf, token, middlewares...) client = &Client{Tg: tgClient, Status: StatusIdle, UserID: userID} w.clients[userID] = client stop, err := Connect(client.Tg, WithBotToken(token)) if err != nil { return nil, err } client.Stop = stop client.Status = StatusBusy w.logger.Debug("started bg client: ", userID) } return client, nil } func (w *StreamWorker) startIdleClientMonitor() { ticker := time.NewTicker(w.cnf.BgBotsCheckInterval) defer ticker.Stop() for { select { case <-ticker.C: w.checkIdleClients() case <-w.ctx.Done(): return } } } func (w *StreamWorker) checkIdleClients() { w.mu.Lock() defer w.mu.Unlock() if w.activeStreams == 0 { for _, client := range w.clients { if client.Status == StatusBusy && client.Stop != nil { client.Stop() client.Stop = nil client.Tg = nil client.Status = StatusIdle w.logger.Debug("stopped bg client: ", client.UserID) } } } }