GlcoinCore/lightningtipbot
1
0
Fork 0
mirror of https://github.com/ChuckNorrison/LightningTipBot.git synced 2026-08-13 12:33:14 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

Merge pull request #404 from LightningTipBot/mutex_states

Browse source 
jobchan
This commit is contained in:
calle 2022-08-31 15:54:05 +02:00 • committed by GitHub
commit 842a037029
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 72 additions and 61 deletions
Download patch file Download diff file Expand all files Collapse all files

133
internal/telegram/generate.go
View file

@ -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,6 +97,21 @@ func (bot *TipBot) confirmGenerateImages(ctx intercept.Context) (intercept.Conte
return ctx, nil
}
var jobChan chan func(workerId int)
var workers = 2
func init() {
jobChan = make(chan func(workerId int), workers)
for i := 0; i < workers; i++ {
go worker(jobChan, i)
}
}
func worker(linkChan chan func(workerId int), workerId int) {
for generatePrompt := range linkChan {
generatePrompt(workerId)
}
}
// generateDalleImages is called by the invoice event when the user has paid
func (bot *TipBot) generateDalleImages(event Event) {
invoiceEvent := event.(*InvoiceEvent)
@ -107,74 +121,71 @@ func (bot *TipBot) generateDalleImages(event Event) {
return
}
bot.trySendMessage(user.Telegram, "🔄 Your images are being generated. Please wait a few moments.")
// we can have only one user using dalle
mutex.Lock("dalle-image-task")
defer mutex.Unlock("dalle-image-task")
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():
bot.dalleRefundUser(user, "")
log.Errorf("[DALLE] ctx done")
var job = func(workerId int) {
// 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-%d] %v", workerId, err.Error())
bot.dalleRefundUser(user,"")
return
// Got a timeout! fail with a timeout error
case <-timeout:
bot.dalleRefundUser(user, "Timeout. Please try again later.")
log.Errorf("[DALLE] 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-%d] %v", workerId, err.Error())
bot.dalleRefundUser(user,"")
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, "")
}
// 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-%d] ctx done", workerId)
return
}
if t.Status == dalle.StatusSucceeded {
log.Infof("[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,"Timeout. Please try again later.")
log.Errorf("[DALLE-%d] timeout", workerId)
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-%d] task succeeded for user %s", worker, 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-%d] %v", worker, err.Error())
}
}
return
} else if t.Status == dalle.StatusRejected {
log.Errorf("[DALLE] rejected: %s", t.ID)
bot.dalleRefundUser(user, "Your prompt has been rejected by OpenAI. Do not use celebrity names, sexual expressions, or any other harmful content as prompt.")
return
} else if t.Status == dalle.StatusRejected {
log.Errorf("[DALLE-%d] rejected: %s", workerId, t.ID)
bot.dalleRefundUser(user,"Your prompt has been rejected by OpenAI. Do not use celebrity names, sexual expressions, or any other harmful content as prompt.")
return
}
log.Debugf("[DALLE-%d] pending for user %s", worker, 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.