long poll exit fix

This commit is contained in:
ycc
2025-05-07 22:53:34 +02:00
parent 8d589505e5
commit f8537aad6d

View File

@ -1,7 +1,6 @@
package helpers
import (
"context"
"errors"
"log"
"os"
@ -194,46 +193,49 @@ func ReadMessage(messageFilename string) ([]string, []string, string, error) {
// CheckForMessages checks for messages on a single server
func LongPollForMessages(storage_path string, jobs []client.RequestsJob, timeout int, longPoll bool) (int, string, error) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
resultChan := make(chan int, len(jobs)) // We'll just send counts here
// Channel to collect results
resultChan := make(chan int, len(jobs))
errChan := make(chan error, len(jobs))
// WaitGroup to sync goroutines
var wg sync.WaitGroup
// Loop through each job (server)
for _, job := range jobs {
wg.Add(1)
go func(job client.RequestsJob) {
defer wg.Done()
for {
select {
case <-ctx.Done():
return
default:
cnt, _, err := CheckForMessages(storage_path, &job, timeout, true)
if err != nil {
continue // Optionally handle/log error
}
if cnt > 0 {
select {
case resultChan <- cnt:
case <-ctx.Done():
}
cancel()
// Long-polling call to the server
cnt, _, err := CheckForMessages(storage_path, &job, timeout, true)
return
}
if err == nil && cnt > 0 {
select {
case resultChan <- cnt:
default:
}
// Close the error channel to notify all goroutines
close(errChan)
}
}(job)
}
// Wait for a result or context timeout
// Close the result channel when all workers are done
go func() {
wg.Wait()
close(resultChan)
}()
// Wait for the first message or all timeouts
select {
case cnt := <-resultChan:
return cnt, "", nil
case <-ctx.Done():
case <-errChan:
// If one fails and exitOnMessage is true
return 0, "", nil
}
}