long poll exit fix
This commit is contained in:
@ -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
|
||||
}
|
||||
|
||||
}
|
||||
|
Reference in New Issue
Block a user