feat(search): add concurrency for page fetching

Introduce concurrency control for Loggly page fetching via the new
`-concurrency` flag. Add OrderedBuffer to ensure ordered delivery of
results. Update search logic to support concurrent requests and
ordered output. Includes tests for OrderedBuffer.

Author: Ajnasz

loggly.go

@@ -30,6 +30,7 @@ const usage = `
     -count            print total event count
     -all              print the entire loggly event instead of just the message
     -maxPages <count> maximum number of pages to query [3]
+    -concurrency <count> number of concurrent page fetchers [3]
     -version          print version information
 
   Operators:
@@ -59,6 +60,7 @@ const usage = `
 // Command options.
 var flags = flag.NewFlagSet("loggly", flag.ExitOnError)
 var count = flags.Bool("count", false, "")
+var concurrency = flags.Int("concurrency", 3, "")
 var versionQuery = flags.Bool("version", false, "")
 var account = flags.String("account", "", "")
 var maxPages = flags.Int("maxPages", 3, "")
@@ -143,10 +145,10 @@ func printRes(res search.Response) {
 	}
 }
 
-func sendQuery(query string, size int, from string, to string, maxPages int) {
+func sendQuery(query string, size int, from string, to string, maxPages int, concurrency int) {
 	doneChan := make(chan error)
 
-	c := search.New(*account, *token)
+	c := search.New(*account, *token).SetConcurrency(concurrency)
 	res, err := c.Query(query).Size(size).From(from).To(to).MaxPage(maxPages).Fetch()
 
 	go func() {
@@ -182,6 +184,12 @@ func warnInvalidFlagPlacement(args []string) {
 	}
 }
 
+func warnHighConcurrency(concurrency int) {
+	if concurrency > 3 {
+		fmt.Fprintf(os.Stderr, " Warning: High concurrency (%d) may lead to rate limiting or temporary blocking by Loggly. Consider reducing the concurrency level.\n", concurrency)
+	}
+}
+
 func main() {
 	flags.Usage = printUsage
 	flags.Parse(os.Args[1:])
@@ -196,12 +204,13 @@ func main() {
 
 	args := flags.Args()
 	warnInvalidFlagPlacement(args)
+	warnHighConcurrency(*concurrency)
 	query := strings.Join(args, " ")
 
 	if *count {
 		execCount(query, *from, *to)
 		return
 	}
 
-	sendQuery(query, *size, *from, *to, *maxPages)
+	sendQuery(query, *size, *from, *to, *maxPages, *concurrency)
 }

orderedbuffer/orderedbuffer.go

@@ -0,0 +1,51 @@
+package orderedbuffer
+
+import (
+	"fmt"
+	"os"
+	"sync"
+)
+
+type OrderedBuffer[T any] struct {
+	responses   map[int]T
+	mu          *sync.RWMutex
+	lastSentIdx int
+	ch          chan T
+}
+
+func NewOrderedBuffer[T any](ch chan T) *OrderedBuffer[T] {
+	fmt.Fprintln(os.Stderr, "Creating new ResponsesStore")
+	var mu sync.RWMutex
+	return &OrderedBuffer[T]{
+		responses:   make(map[int]T),
+		ch:          ch,
+		lastSentIdx: -1,
+		mu:          &mu,
+	}
+}
+
+func (s *OrderedBuffer[T]) send() {
+	s.mu.Lock()
+	newIdx := s.lastSentIdx + 1
+	fmt.Fprintln(os.Stderr, "Checking for page", newIdx)
+	if resp, ok := s.responses[newIdx]; ok {
+		fmt.Fprintln(os.Stderr, "Sending page", newIdx)
+		s.ch <- resp
+		s.lastSentIdx = newIdx
+		delete(s.responses, newIdx)
+		s.mu.Unlock()
+		s.send()
+		return
+	} else {
+		fmt.Fprintln(os.Stderr, "Page", newIdx, "not ready yet")
+	}
+	s.mu.Unlock()
+}
+
+func (s *OrderedBuffer[T]) Store(i int, r T) {
+	fmt.Fprintf(os.Stderr, "Storing page %d\n", i)
+	s.mu.Lock()
+	s.responses[i] = r
+	s.mu.Unlock()
+	s.send()
+}

orderedbuffer/orderedbuffer_test.go

@@ -0,0 +1,58 @@
+package orderedbuffer
+
+import (
+	"testing"
+	"time"
+)
+
+func TestOrderedBufferOrderedDelivery(t *testing.T) {
+	ch := make(chan int, 3)
+	buf := NewOrderedBuffer(ch)
+
+	buf.Store(0, 10)
+	buf.Store(1, 20)
+	buf.Store(2, 30)
+
+	got := []int{<-ch, <-ch, <-ch}
+	want := []int{10, 20, 30}
+	for i := range want {
+		if got[i] != want[i] {
+			t.Errorf("expected %d at index %d, got %d", want[i], i, got[i])
+		}
+	}
+}
+
+func TestOrderedBufferOutOfOrder(t *testing.T) {
+	ch := make(chan string, 3)
+	buf := NewOrderedBuffer(ch)
+
+	buf.Store(2, "c")
+	buf.Store(0, "a")
+	buf.Store(1, "b")
+
+	got := []string{<-ch, <-ch, <-ch}
+	want := []string{"a", "b", "c"}
+	for i := range want {
+		if got[i] != want[i] {
+			t.Errorf("expected %s at index %d, got %s", want[i], i, got[i])
+		}
+	}
+}
+
+func TestOrderedBufferConcurrent(t *testing.T) {
+	ch := make(chan int, 3)
+	buf := NewOrderedBuffer(ch)
+
+	go buf.Store(1, 100)
+	go buf.Store(0, 50)
+	go buf.Store(2, 150)
+
+	time.Sleep(100 * time.Millisecond)
+	got := []int{<-ch, <-ch, <-ch}
+	want := []int{50, 100, 150}
+	for i := range want {
+		if got[i] != want[i] {
+			t.Errorf("expected %d at index %d, got %d", want[i], i, got[i])
+		}
+	}
+}

search/search.go

@@ -2,28 +2,32 @@ package search
 
 import (
 	"fmt"
-	"io/ioutil"
+	"io"
 	"net/http"
 	"net/url"
 	"strconv"
+	"sync"
+	"sync/atomic"
 
+	"github.com/Ajnasz/go-loggly-cli/orderedbuffer"
 	"github.com/bitly/go-simplejson"
 )
 
 // Client Loggly search client with user credentials, loggly
 // does not seem to support tokens right now.
 type Client struct {
-	Token    string
-	Account  string
-	Endpoint string
+	Token       string
+	Account     string
+	Endpoint    string
+	concurrency int
 }
 
 // Response Search response with total events, page number
 // and the events array.
 type Response struct {
 	Total  int64
 	Page   int64
-	Events []interface{}
+	Events []any
 }
 
 // Query builder struct
@@ -60,6 +64,14 @@ func New(account string, token string) *Client {
 	return c
 }
 
+func (c *Client) SetConcurrency(n int) *Client {
+	if n < 1 {
+		n = 1
+	}
+	c.concurrency = n
+	return c
+}
+
 // URL Return the base api url.
 func (c *Client) URL() string {
 	return fmt.Sprintf("https://%s.%s", c.Account, c.Endpoint)
@@ -89,10 +101,11 @@ func (c *Client) GetJSON(path string) (j *simplejson.Json, err error) {
 	defer res.Body.Close()
 
 	if res.StatusCode >= 400 {
-		return nil, fmt.Errorf("go-loggly-search: %q", res.Status)
+		body, _ := io.ReadAll(res.Body)
+		return nil, fmt.Errorf("go-loggly-search: %q, %s", res.Status, body)
 	}
 
-	body, err := ioutil.ReadAll(res.Body)
+	body, err := io.ReadAll(res.Body)
 
 	if err != nil {
 		return
@@ -187,7 +200,12 @@ func (q *Query) Fetch() (chan Response, chan error) {
 	resChan := make(chan Response)
 	errChan := make(chan error)
 
-	page := 0
+	concurrent := min(q.maxPages, q.client.concurrency)
+
+	pool := make(chan struct{}, concurrent)
+
+	var page atomic.Int64
+	page.Store(-1)
 	go func() {
 		defer close(resChan)
 		defer close(errChan)
@@ -198,26 +216,45 @@ func (q *Query) Fetch() (chan Response, chan error) {
 			return
 		}
 
+		var wg sync.WaitGroup
+		var hasMore atomic.Bool
+		hasMore.Store(true)
+		var lastPrintedPage atomic.Int32
+		lastPrintedPage.Store(-1)
+		responsesStore := orderedbuffer.NewOrderedBuffer(resChan)
 		for {
-			res, err := q.client.Search(j, page)
-
-			if err != nil {
-				errChan <- err
-				return
+			pool <- struct{}{}
+			wg.Add(1)
+			go func(page int) {
+				defer wg.Done()
+				defer func() { <-pool }()
+
+				res, err := q.client.Search(j, page)
+				if err != nil {
+					errChan <- err
+					hasMore.Store(false)
+					return
+				}
+
+				if res != nil {
+					// resChan <- *res
+					responsesStore.Store(page, *res)
+
+					if len(res.Events) < q.size {
+						hasMore.Store(false)
+						return
+					}
+				} else {
+					hasMore.Store(false)
+					return
+				}
+			}(int(page.Add(1)))
+			if int(page.Load()) >= q.maxPages || !hasMore.Load() {
+				break
 			}
-
-			resChan <- *res
-
-			if page+1 == q.maxPages {
-				return
-			}
-
-			if len(res.Events) < q.size {
-				return
-			}
-
-			page++
 		}
+
+		wg.Wait()
 	}()
 
 	return resChan, errChan