Các Mẫu Goroutine Pool Trong Go: errgroup & Worker Pool

package concurrency

import (
	"context"
	"fmt"
	"log/slog"
	"time"

	"golang.org/x/sync/errgroup"
)

// ProductItem định nghĩa payload của từng đơn vị xử lý
type ProductItem struct {
	ID    int64
	SKU   string
	Price float64
}

// BatchPriceUpdater quản lý tác vụ cập nhật giá sản phẩm hàng loạt
type BatchPriceUpdater struct {
	concurrencyLimit int
	timeout          time.Duration
}

// NewBatchPriceUpdater khởi tạo updater với giới hạn worker đồng thời
func NewBatchPriceUpdater(limit int, timeout time.Duration) *BatchPriceUpdater {
	if limit <= 0 {
		limit = 10
	}
	return &BatchPriceUpdater{
		concurrencyLimit: limit,
		timeout:          timeout,
	}
}

// UpdatePrices thực thi cập nhật danh sách sản phẩm với errgroup
func (u *BatchPriceUpdater) UpdatePrices(parentCtx context.Context, items []ProductItem) error {
	ctx, cancel := context.WithTimeout(parentCtx, u.timeout)
	defer cancel()

	// Khởi tạo errgroup gắn kèm context hủy bỏ tự động khi có lỗi phát sinh
	g, groupCtx := errgroup.WithContext(ctx)
	g.SetLimit(u.concurrencyLimit)

	for _, item := range items {
		item := item // Bảo vệ biến lặp (local closure scope)

		// g.Go sẽ block caller goroutine khi số lượng worker đạt concurrencyLimit
		g.Go(func() error {
			// Kiểm tra sớm tín hiệu hủy từ các goroutine anh em
			select {
			case <-groupCtx.Done():
				return groupCtx.Err()
			default:
			}

			if err := u.processItem(groupCtx, item); err != nil {
				slog.Error("Cập nhật sản phẩm thất bại", "sku", item.SKU, "error", err)
				// Lỗi đầu tiên trả về sẽ kích hoạt hủy groupCtx cho tất cả tác vụ còn lại
				return fmt.Errorf("cập nhật sku %s thất bại: %w", item.SKU, err)
			}
			return nil
		})
	}

	// Chờ đợi tất cả goroutines hoàn thành hoặc gặp lỗi đầu tiên
	if err := g.Wait(); err != nil {
		return fmt.Errorf("tiến trình xử lý theo lô thất bại: %w", err)
	}

	slog.Info("Cập nhật toàn bộ lô sản phẩm thành công", "total", len(items))
	return nil
}

func (u *BatchPriceUpdater) processItem(ctx context.Context, item ProductItem) error {
	// Giả lập I/O database call
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-time.After(25 * time.Millisecond):
		return nil
	}
}

package concurrency

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"sync"
	"time"

	"golang.org/x/sync/semaphore"
)

var (
	ErrQueueSaturated = errors.New("hệ thống đang bận, vượt ngưỡng concurrency cho phép")
)

// ResourceController quản lý cấp phát tài nguyên theo trọng số
type ResourceController struct {
	sem        *semaphore.Weighted
	maxTickets int64
	wg         sync.WaitGroup
}

func NewResourceController(totalCapacity int64) *ResourceController {
	return &ResourceController{
		sem:        semaphore.NewWeighted(totalCapacity),
		maxTickets: totalCapacity,
	}
}

// ExecuteWithPriority thực thi tác vụ với số lượng ticket chỉ định
func (rc *ResourceController) ExecuteWithPriority(ctx context.Context, tickets int64, task func(context.Context) error) error {
	// Non-blocking try: Kiểm tra nếu không còn đủ slot thì reject ngay (Fast Backpressure)
	if !rc.sem.TryAcquire(tickets) {
		return ErrQueueSaturated
	}

	rc.wg.Add(1)
	go func() {
		defer rc.wg.Done()
		defer rc.sem.Release(tickets)

		// context.WithoutCancel giúp bảo toàn tác vụ đang ghi dở khi client ngắt kết nối bất chợt
		execCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 30*time.Second)
		defer cancel()

		if err := task(execCtx); err != nil {
			slog.Error("Thực thi task thất bại", "tickets", tickets, "error", err)
		}
	}()

	return nil
}

// GracefulWait đóng nhận tác vụ mới và đợi các tác vụ hiện hữu hoàn tất
func (rc *ResourceController) GracefulWait(ctx context.Context) error {
	done := make(chan struct{})
	go func() {
		rc.wg.Wait()
		close(done)
	}()

	select {
	case <-done:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

package workerpool

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"log/slog"
	"sync"
	"sync/atomic"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
)

var (
	ErrPoolClosed   = errors.New("worker pool đã dừng nhận tác vụ")
	ErrQueueFull    = errors.New("hàng đợi xử lý đã đầy (backpressure reject)")
	ErrDrainTimeout = errors.New("quá thời gian chờ giải phóng worker pool")
)

// Các metrics giám sát hệ thống thời gian thực
var (
	metricJobsSubmitted = promauto.NewCounterVec(prometheus.CounterOpts{
		Name: "goroutine_pool_jobs_submitted_total",
		Help: "Tổng số lượng tác vụ được đưa vào worker pool",
	}, []string{"pool_name"})

	metricJobsDropped = promauto.NewCounterVec(prometheus.CounterOpts{
		Name: "goroutine_pool_jobs_dropped_total",
		Help: "Tổng số lượng tác vụ bị từ chối do quá tải hàng đợi",
	}, []string{"pool_name", "reason"})

	metricQueueDepth = promauto.NewGaugeVec(prometheus.GaugeOpts{
		Name: "goroutine_pool_queue_depth",
		Help: "Độ sâu hàng đợi hiện tại của worker pool",
	}, []string{"pool_name"})

	metricActiveWorkers = promauto.NewGaugeVec(prometheus.GaugeOpts{
		Name: "goroutine_pool_active_workers",
		Help: "Số lượng worker hiện đang bận xử lý công việc",
	}, []string{"pool_name"})

	metricJobDuration = promauto.NewHistogramVec(prometheus.HistogramOpts{
		Name:    "goroutine_pool_job_duration_seconds",
		Help:    "Thời gian xử lý từng tác vụ trong pool",
		Buckets: []float64{0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5},
	}, []string{"pool_name", "status"})
)

// BufferPool tái sử dụng bộ nhớ đệm bytes.Buffer nhằm triệt tiêu cấp phát heap
var bufferPool = sync.Pool{
	New: func() any {
		return bytes.NewBuffer(make([]byte, 0, 4096))
	},
}

// JobHandler định nghĩa hàm thực thi nghiệp vụ cho từng task
type JobHandler[T any] func(ctx context.Context, payload T, buf *bytes.Buffer) error

// ProductionPool quản lý hàng đợi và tập hợp worker
type ProductionPool[T any] struct {
	name         string
	workerCount  int
	queueCap     int
	jobQueue     chan T
	handler      JobHandler[T]
	wg           sync.WaitGroup
	isClosed     atomic.Bool
	activeCount  atomic.Int64
	closedChan   chan struct{}
}

// NewProductionPool khởi tạo pool với kích thước hàng đợi và số worker cố định
func NewProductionPool[T any](name string, workers int, queueCap int, handler JobHandler[T]) *ProductionPool[T] {
	if workers <= 0 {
		workers = 10
	}
	if queueCap <= 0 {
		queueCap = 100
	}

	return &ProductionPool[T]{
		name:        name,
		workerCount: workers,
		queueCap:    queueCap,
		jobQueue:    make(chan T, queueCap),
		handler:     handler,
		closedChan:  make(chan struct{}),
	}
}

// Start kích hoạt toàn bộ worker goroutines
func (p *ProductionPool[T]) Start(ctx context.Context) {
	for i := 0; i < p.workerCount; i++ {
		p.wg.Add(1)
		go p.workerLoop(ctx, i)
	}
	slog.Info("Worker pool đã sẵn sàng hoạt động", "pool", p.name, "workers", p.workerCount, "capacity", p.queueCap)
}

func (p *ProductionPool[T]) workerLoop(ctx context.Context, workerID int) {
	defer p.wg.Done()

	for {
		// Cập nhật metric độ sâu hàng đợi định kỳ
		metricQueueDepth.WithLabelValues(p.name).Set(float64(len(p.jobQueue)))

		select {
		case <-ctx.Done():
			// Tín hiệu hủy từ context cha
			return

		case <-p.closedChan:
			// Chuyển sang chế độ vét cạn hàng đợi (Drain Mode)
			for job := range p.jobQueue {
				p.executeJob(context.Background(), job)
			}
			return

		case job, ok := <-p.jobQueue:
			if !ok {
				// Channel đã đóng hoàn toàn
				return
			}
			p.executeJob(ctx, job)
		}
	}
}

func (p *ProductionPool[T]) executeJob(ctx context.Context, job T) {
	p.activeCount.Add(1)
	metricActiveWorkers.WithLabelValues(p.name).Set(float64(p.activeCount.Load()))
	start := time.Now()

	// Tái sử dụng byte buffer từ sync.Pool
	buf := bufferPool.Get().(*bytes.Buffer)
	buf.Reset()
	defer func() {
		bufferPool.Put(buf)
		p.activeCount.Add(-1)
		metricActiveWorkers.WithLabelValues(p.name).Set(float64(p.activeCount.Load()))
	}()

	status := "success"
	if err := p.handler(ctx, job, buf); err != nil {
		status = "error"
		slog.Error("Xử lý job trong pool thất bại", "pool", p.name, "error", err)
	}

	duration := time.Since(start).Seconds()
	metricJobDuration.WithLabelValues(p.name, status).Observe(duration)
}

// Submit đưa tác vụ vào hàng đợi với cơ chế chặn an toàn (blocking backpressure)
func (p *ProductionPool[T]) Submit(ctx context.Context, job T) error {
	if p.isClosed.Load() {
		return ErrPoolClosed
	}

	select {
	case <-ctx.Done():
		return ctx.Err()
	case p.jobQueue <- job:
		metricJobsSubmitted.WithLabelValues(p.name).Inc()
		return nil
	}
}

// SubmitOrDrop cố gắng đưa tác vụ vào pool; nếu đầy thì từ chối ngay lập tức (Non-blocking)
func (p *ProductionPool[T]) SubmitOrDrop(job T) error {
	if p.isClosed.Load() {
		return ErrPoolClosed
	}

	select {
	case p.jobQueue <- job:
		metricJobsSubmitted.WithLabelValues(p.name).Inc()
		return nil
	default:
		metricJobsDropped.WithLabelValues(p.name, "queue_full").Inc()
		return ErrQueueFull
	}
}

// DrainAndShutdown dừng tiếp nhận tác vụ mới và chờ xử lý hết hàng đợi trong thời hạn timeout
func (p *ProductionPool[T]) DrainAndShutdown(timeout time.Duration) error {
	if !p.isClosed.CompareAndSwap(false, true) {
		return nil // Đã được gọi trước đó
	}

	slog.Info("Bắt đầu quy trình Graceful Shutdown cho worker pool", "pool", p.name)
	close(p.closedChan) // Báo hiệu cho các worker chuyển sang vét hàng đợi

	done := make(chan struct{})
	go func() {
		p.wg.Wait()
		close(done)
	}()

	select {
	case <-done:
		slog.Info("Worker pool đã được giải phóng toàn diện", "pool", p.name)
		return nil
	case <-time.After(timeout):
		return fmt.Errorf("%w: còn lại %d jobs chưa kịp xử lý", ErrDrainTimeout, len(p.jobQueue))
	}
}

package workerpool_test

import (
	"context"
	"testing"
	"time"

	"bytes"
	"d:/myproject/learn/workerpool" // Điều chỉnh theo go module thực tế
	"go.uber.org/goleak"
)

func TestWorkerPool_NoGoroutineLeak(t *testing.T) {
	// goleak.VerifyNone kiểm tra toàn bộ goroutine còn sống sau khi test kết thúc
	defer goleak.VerifyNone(t)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	processedCount := 0
	handler := func(ctx context.Context, item int, buf *bytes.Buffer) error {
		buf.WriteString("ok")
		processedCount++
		return nil
	}

	pool := workerpool.NewProductionPool[int]("test-pool", 5, 20, handler)
	pool.Start(ctx)

	for i := 0; i < 50; i++ {
		if err := pool.Submit(ctx, i); err != nil {
			t.Fatalf("Không thể submit job %d: %v", i, err)
		}
	}

	// Thực hiện Graceful Drain
	if err := pool.DrainAndShutdown(3 * time.Second); err != nil {
		t.Fatalf("Drain thất bại: %v", err)
	}

	if processedCount != 50 {
		t.Errorf("Kỳ vọng xử lý 50 jobs, thực tế: %d", processedCount)
	}
}