Move flushall to queue.Manager and add to testlogger

Author: zeripath

integrations/testlogger.go

@@ -5,6 +5,7 @@
 package integrations
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"os"
@@ -100,13 +101,8 @@ func PrintCurrentTest(t testing.TB, skip ...int) func() {
 	}
 	writerCloser.setT(&t)
 	return func() {
-		for i := 0; i < 5; i++ {
-			qs := queue.GetManager().ManagedQueues()
-			for _, q := range qs {
-				if err := q.Flush(10 * time.Second); err != nil {
-					t.Errorf("Flushing queue %s failed with error %v", q.Name, err)
-				}
-			}
+		if err := queue.GetManager().FlushAll(context.Background(), 20*time.Second); err != nil {
+			t.Errorf("Flushing queues failed with error %v", err)
 		}
 		_ = writerCloser.Close()
 	}

modules/queue/manager.go

@@ -44,8 +44,11 @@ type ManagedQueue struct {
 
 // Flushable represents a pool or queue that is flushable
 type Flushable interface {
-	// Flush will add a flush worker to the pool
+	// Flush will add a flush worker to the pool - the worker should be autoregistered with the manager
 	Flush(time.Duration) error
+	// FlushWithContext is very similar to Flush
+	// NB: The worker will not be registered with the manager.
+	FlushWithContext(ctx context.Context) error
 	// IsEmpty will return if the managed pool is empty and has no work
 	IsEmpty() bool
 }
@@ -147,6 +150,64 @@ func (m *Manager) GetManagedQueue(qid int64) *ManagedQueue {
 	return m.Queues[qid]
 }
 
+// FlushAll flushes all the flushable queues attached to this manager
+func (m *Manager) FlushAll(baseCtx context.Context, timeout time.Duration) error {
+	var ctx context.Context
+	var cancel context.CancelFunc
+	start := time.Now()
+	end := start
+	hasTimeout := false
+	if timeout > 0 {
+		ctx, cancel = context.WithTimeout(baseCtx, timeout)
+		end = start.Add(timeout)
+		hasTimeout = true
+	} else {
+		ctx, cancel = context.WithCancel(baseCtx)
+	}
+	defer cancel()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		default:
+		}
+		mqs := m.ManagedQueues()
+		wg := sync.WaitGroup{}
+		wg.Add(len(mqs))
+		allEmpty := true
+		for _, mq := range mqs {
+			if mq.IsEmpty() {
+				wg.Done()
+				continue
+			}
+			allEmpty = false
+			if flushable, ok := mq.Managed.(Flushable); ok {
+				go func() {
+					localCtx, localCancel := context.WithCancel(ctx)
+					pid := mq.RegisterWorkers(1, start, hasTimeout, end, localCancel, true)
+					err := flushable.FlushWithContext(localCtx)
+					if err != nil && err != ctx.Err() {
+						cancel()
+					}
+					mq.CancelWorkers(pid)
+					localCancel()
+					wg.Done()
+				}()
+			} else {
+				wg.Done()
+			}
+
+		}
+		if allEmpty {
+			break
+		}
+		wg.Wait()
+	}
+	return nil
+
+}
+
 // ManagedQueues returns the managed queues
 func (m *Manager) ManagedQueues() []*ManagedQueue {
 	m.mutex.Lock()

routers/private/manager.go

@@ -5,10 +5,7 @@
 package private
 
 import (
-	"context"
 	"net/http"
-	"sync"
-	"time"
 
 	"code.gitea.io/gitea/modules/graceful"
 	"code.gitea.io/gitea/modules/log"
@@ -20,28 +17,15 @@ import (
 
 // FlushQueues flushes all the Queues
 func FlushQueues(ctx *macaron.Context, opts private.FlushOptions) {
-	var flushCtx context.Context
-	var cancel context.CancelFunc
-	start := time.Now()
-	end := start
-	hasTimeout := false
 
 	baseCtx := ctx.Req.Request.Context()
 	if opts.NonBlocking {
 		baseCtx = graceful.GetManager().HammerContext()
 	}
 
-	if opts.Timeout > 0 {
-		flushCtx, cancel = context.WithTimeout(baseCtx, opts.Timeout)
-		end = start.Add(opts.Timeout)
-		hasTimeout = true
-	} else {
-		flushCtx, cancel = context.WithCancel(baseCtx)
-	}
-
 	if opts.NonBlocking {
 		go func() {
-			err := doFlush(flushCtx, cancel, start, end, hasTimeout)
+			err := queue.GetManager().FlushAll(baseCtx, opts.Timeout)
 			if err != nil {
 				log.Error("Flushing request timed-out with error: %v", err)
 			}
@@ -51,54 +35,11 @@ func FlushQueues(ctx *macaron.Context, opts private.FlushOptions) {
 		})
 		return
 	}
-	err := doFlush(flushCtx, cancel, start, end, hasTimeout)
+	err := queue.GetManager().FlushAll(baseCtx, opts.Timeout)
 	if err != nil {
 		ctx.JSON(http.StatusRequestTimeout, map[string]interface{}{
-			"err": flushCtx.Err().Error(),
+			"err": err,
 		})
 	}
 	ctx.PlainText(http.StatusOK, []byte("success"))
 }
-
-func doFlush(ctx context.Context, cancel context.CancelFunc, start, end time.Time, hasTimeout bool) error {
-	defer cancel()
-	for {
-		select {
-		case <-ctx.Done():
-			return ctx.Err()
-		default:
-		}
-		mqs := queue.GetManager().ManagedQueues()
-		wg := sync.WaitGroup{}
-		wg.Add(len(mqs))
-		allEmpty := true
-		for _, mq := range mqs {
-			if mq.IsEmpty() {
-				wg.Done()
-				continue
-			}
-			allEmpty = false
-			if pool, ok := mq.Managed.(queue.WorkerPool); ok {
-				go func() {
-					localCtx, localCancel := context.WithCancel(ctx)
-					pid := mq.RegisterWorkers(1, start, hasTimeout, end, localCancel, true)
-					err := pool.FlushWithContext(localCtx)
-					if err != nil && err != ctx.Err() {
-						cancel()
-					}
-					mq.CancelWorkers(pid)
-					localCancel()
-					wg.Done()
-				}()
-			} else {
-				wg.Done()
-			}
-
-		}
-		if allEmpty {
-			break
-		}
-		wg.Wait()
-	}
-	return nil
-}