diff --git a/internal/telegram/generate.go b/internal/telegram/generate.go index ca7d054..2523a6b 100644 --- a/internal/telegram/generate.go +++ b/internal/telegram/generate.go @@ -13,7 +13,6 @@ import ( "github.com/LightningTipBot/LightningTipBot/internal/dalle" "github.com/LightningTipBot/LightningTipBot/internal/lnbits" "github.com/LightningTipBot/LightningTipBot/internal/runtime" - "github.com/LightningTipBot/LightningTipBot/internal/runtime/mutex" "github.com/LightningTipBot/LightningTipBot/internal/telegram/intercept" log "github.com/sirupsen/logrus" "github.com/skip2/go-qrcode" @@ -98,7 +97,18 @@ func (bot *TipBot) confirmGenerateImages(ctx intercept.Context) (intercept.Conte return ctx, nil } -var mutexStates = map[int]bool{0: false, 1: false} +var jobChan chan func() + +func init() { + jobChan = make(chan func(), 2) + go worker(jobChan) + go worker(jobChan) +} +func worker(linkChan chan func()) { + for generatePrompt := range linkChan { + generatePrompt() + } +} // generateDalleImages is called by the invoice event when the user has paid func (bot *TipBot) generateDalleImages(event Event) { @@ -109,86 +119,71 @@ func (bot *TipBot) generateDalleImages(event Event) { return } bot.trySendMessage(user.Telegram, "🔄 Your images are being generated. Please wait a few moments.") - locker := -1 - for i := 0; i < 2; i++ { - if mutexStates[i] == false { - mutexStates[i] = true - locker = i - mutex.Lock(fmt.Sprintf("dalle-image-task-%d", i)) - break - } - if i == len(mutexStates)-1 { - time.Sleep(time.Second * 1) - i = 0 - } - } - if locker > 0 { - defer mutex.Unlock(fmt.Sprintf("dalle-image-task-%d", locker)) - } - time.Sleep(time.Second * 1) - - // create the client with the bearer token api key - dalleClient, err := dalle.NewHTTPClient(internal.Configuration.Generate.DalleKey) - // handle err - if err != nil { - log.Errorf("[NewHTTPClient] %v", err.Error()) - bot.dalleRefundUser(user) - return - } - - ctx, cancel := context.WithTimeout(context.Background(), time.Minute*10) - defer cancel() - // generate a task to create an image with a prompt - task, err := dalleClient.Generate(ctx, invoiceEvent.CallbackData) - if err != nil { - log.Errorf("[Generate] %v", err.Error()) - bot.dalleRefundUser(user) - return - } - // poll the task.ID until status is succeeded - var t *dalle.Task - timeout := time.After(5 * time.Minute) - ticker := time.Tick(5 * time.Second) - // Keep trying until we're timed out or get a result/error - for { - select { - case <-ctx.Done(): + var job = func() { + // create the client with the bearer token api key + dalleClient, err := dalle.NewHTTPClient(internal.Configuration.Generate.DalleKey) + // handle err + if err != nil { + log.Errorf("[NewHTTPClient] %v", err.Error()) bot.dalleRefundUser(user) - log.Errorf("[DALLE] ctx done") return - // Got a timeout! fail with a timeout error - case <-timeout: + } + + ctx, cancel := context.WithTimeout(context.Background(), time.Minute*10) + defer cancel() + // generate a task to create an image with a prompt + task, err := dalleClient.Generate(ctx, invoiceEvent.CallbackData) + if err != nil { + log.Errorf("[Generate] %v", err.Error()) bot.dalleRefundUser(user) - log.Errorf("[DALLE] timeout") return - // Got a tick, we should check on checkSomething() - case <-ticker: - t, err = dalleClient.GetTask(ctx, task.ID) - // handle err - if err != nil { - log.Errorf("[GetTask] %v", err.Error()) + } + // poll the task.ID until status is succeeded + var t *dalle.Task + timeout := time.After(5 * time.Minute) + ticker := time.Tick(5 * time.Second) + // Keep trying until we're timed out or get a result/error + for { + select { + case <-ctx.Done(): bot.dalleRefundUser(user) + log.Errorf("[DALLE] ctx done") return - } - if t.Status == dalle.StatusSucceeded { - fmt.Printf("[DALLE] task succeeded for user %s", GetUserStr(user.Telegram)) - // download the first generated image - for _, data := range t.Generations.Data { - err = bot.downloadAndSendImages(ctx, dalleClient, data, invoiceEvent) - if err != nil { - log.Errorf("[downloadAndSendImages] %v", err.Error()) - } + // Got a timeout! fail with a timeout error + case <-timeout: + bot.dalleRefundUser(user) + log.Errorf("[DALLE] timeout") + return + // Got a tick, we should check on checkSomething() + case <-ticker: + t, err = dalleClient.GetTask(ctx, task.ID) + // handle err + if err != nil { + log.Errorf("[GetTask] %v", err.Error()) + bot.dalleRefundUser(user) + return } - return + if t.Status == dalle.StatusSucceeded { + fmt.Printf("[DALLE] task succeeded for user %s", GetUserStr(user.Telegram)) + // download the first generated image + for _, data := range t.Generations.Data { + err = bot.downloadAndSendImages(ctx, dalleClient, data, invoiceEvent) + if err != nil { + log.Errorf("[downloadAndSendImages] %v", err.Error()) + } + } + return - } else if t.Status == dalle.StatusRejected { - log.Errorf("[DALLE] rejected: %s", t.ID) - bot.dalleRefundUser(user) - return + } else if t.Status == dalle.StatusRejected { + log.Errorf("[DALLE] rejected: %s", t.ID) + bot.dalleRefundUser(user) + return + } + log.Debugf("[DALLE] pending for user %s", GetUserStr(user.Telegram)) } - log.Debugf("[DALLE] pending for user %s", GetUserStr(user.Telegram)) } } + jobChan <- job } // downloadAndSendImages will download dalle images and send them to the payer.