teldrive/internal/tgc/workers.go

package tgc

import (
	"context"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/divyam234/teldrive/internal/config"
	"github.com/divyam234/teldrive/internal/kv"
	"github.com/divyam234/teldrive/internal/logging"
	"github.com/gotd/td/telegram"
	"go.uber.org/zap"
)

type UploadWorker struct {
	mu      sync.Mutex
	bots    map[int64][]string
	currIdx map[int64]int
}

func (w *UploadWorker) Set(bots []string, channelId int64) {
	w.mu.Lock()
	defer w.mu.Unlock()
	_, ok := w.bots[channelId]
	if !ok {
		w.bots = make(map[int64][]string)
		w.currIdx = make(map[int64]int)
		w.bots[channelId] = bots
		w.currIdx[channelId] = 0
	}
}

func (w *UploadWorker) Next(channelId int64) (string, int) {
	w.mu.Lock()
	defer w.mu.Unlock()
	index := w.currIdx[channelId]
	w.currIdx[channelId] = (index + 1) % len(w.bots[channelId])
	return w.bots[channelId][index], index
}

func NewUploadWorker() *UploadWorker {
	return &UploadWorker{}
}

type Client struct {
	Tg          *telegram.Client
	Stop        StopFunc
	Status      string
	UserId      string
	lastUsed    time.Time
	connections int
}

type StreamWorker struct {
	mu          sync.Mutex
	clients     map[string]*Client
	currIdx     map[int64]int
	channelBots map[int64][]string
	cnf         *config.TGConfig
	kv          kv.KV
	ctx         context.Context
	logger      *zap.SugaredLogger
}

func (w *StreamWorker) Set(bots []string, channelId int64) {

	w.mu.Lock()
	defer w.mu.Unlock()
	_, ok := w.channelBots[channelId]
	if !ok {
		w.channelBots[channelId] = bots
		w.currIdx[channelId] = 0
	}

}

func (w *StreamWorker) Next(channelId int64) (*Client, int, error) {
	w.mu.Lock()
	defer w.mu.Unlock()
	index := w.currIdx[channelId]
	token := w.channelBots[channelId][index]
	userId := strings.Split(token, ":")[0]
	client, ok := w.clients[userId]
	if !ok || (client.Status == "idle" && client.Stop == nil) {
		middlewares := Middlewares(w.cnf, 5)
		tgClient, _ := BotClient(w.ctx, w.kv, w.cnf, token, middlewares...)
		client = &Client{Tg: tgClient, Status: "idle", UserId: userId}
		w.clients[userId] = client
		stop, err := Connect(client.Tg, WithBotToken(token))
		if err != nil {
			return nil, 0, err
		}
		client.Stop = stop
		w.logger.Debug("started bg client: ", client.UserId)
	}
	w.currIdx[channelId] = (index + 1) % len(w.channelBots[channelId])
	client.lastUsed = time.Now()
	if client.connections == 0 {
		client.Status = "serving"
	}
	client.connections++
	return client, index, nil
}

func (w *StreamWorker) Release(client *Client) {
	w.mu.Lock()
	defer w.mu.Unlock()
	client.connections--
	if client.connections == 0 {
		client.Status = "running"
	}
}

func (w *StreamWorker) UserWorker(session string, userId int64) (*Client, error) {
	w.mu.Lock()
	defer w.mu.Unlock()

	id := strconv.FormatInt(userId, 10)
	client, ok := w.clients[id]
	if !ok || (client.Status == "idle" && client.Stop == nil) {
		middlewares := Middlewares(w.cnf, 5)
		tgClient, _ := AuthClient(w.ctx, w.cnf, session, middlewares...)
		client = &Client{Tg: tgClient, Status: "idle", UserId: id}
		w.clients[id] = client
		stop, err := Connect(client.Tg, WithContext(w.ctx))
		if err != nil {
			return nil, err
		}
		client.Stop = stop
		w.logger.Debug("started bg client: ", client.UserId)
	}
	client.lastUsed = time.Now()
	if client.connections == 0 {
		client.Status = "serving"
	}
	client.connections++
	return client, nil
}

func (w *StreamWorker) startIdleClientMonitor() {
	ticker := time.NewTicker(w.cnf.BgBotsCheckInterval)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			w.mu.Lock()
			for _, client := range w.clients {
				if client.Status == "running" && time.Since(client.lastUsed) > w.cnf.BgBotsTimeout {
					if client.Stop != nil {
						client.Stop()
						client.Stop = nil
						client.Status = "idle"
						w.logger.Debug("stopped bg client: ", client.UserId)
					}
				}
			}
			w.mu.Unlock()
		case <-w.ctx.Done():
			return
		}
	}

}

func NewStreamWorker(ctx context.Context) func(cnf *config.Config, kv kv.KV) *StreamWorker {
	return func(cnf *config.Config, kv kv.KV) *StreamWorker {
		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: logging.FromContext(ctx)}
		go worker.startIdleClientMonitor()
		return worker
	}
}