Golang concurrency pool

main.go

package main

import (
	"context"
	"fmt"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

// ConcPool implementation with go channels
type ConcPool struct {
	ch chan struct{}
}

func (cp *ConcPool) Release() {
	cp.ch <- struct{}{}
}

func (cp *ConcPool) GetLease() {
	<-cp.ch
}

func InitPool(size int) *ConcPool {
	cp := &ConcPool{ch: make(chan struct{}, size)}
	for i := 0; i < size; i++ {
		cp.Release()
	}
	return cp
}

func (p *ConcPool) GetLeaseWithCtx(ctx context.Context) error {
	select {
	case <-ctx.Done():
		return ctx.Err()
	default:
	}

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

// TestPool tests the Concurrency Pool implementation
func TestPool(t *testing.T) {
	// Initialize a pool with 3 slots
	pool := InitPool(3)

	// Test: Verify concurrent access
	t.Run("Concurrent Access", func(t *testing.T) {
		var wg sync.WaitGroup
		var successCount, errorCount atomic.Uint32

		// Simulate 10 goroutines trying to acquire slots
		for i := 0; i < 10; i++ {
			wg.Add(1)
			go func(n int) {
				defer wg.Done()

				ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
				defer cancel()

				err := pool.GetLeaseWithCtx(ctx)
				if err != nil {
					errorCount.Add(1)
					fmt.Printf("Goroutine %d: Failed to acquire a slot: %v\n", n, err)
				} else {
					successCount.Add(1)
					fmt.Printf("Goroutine %d: Successfully acquired a slot\n", n)
					time.Sleep(1 * time.Second) // Simulate work
					pool.Release()
				}
			}(i)
		}

		wg.Wait()

		// Verify that only 3 goroutines succeeded at a time
		if successCount.Load() > 3 {
			t.Errorf("Expected at most 3 successful slots, but got %d", successCount.Load())
		} else {
			fmt.Printf("Test 3: Successfully limited concurrency to 3 (success: %d, errors: %d)\n", successCount.Load(), errorCount.Load())
		}
	})
}

func main() {
	// Run the tests
	testing.Main(func(pat, str string) (bool, error) { return true, nil }, []testing.InternalTest{
		{Name: "TestPool", F: TestPool},
	}, nil, nil)
}

/*
➜  go run main.go
Goroutine 9: Successfully acquired lease
Goroutine 5: Successfully acquired lease
Goroutine 0: Successfully acquired lease
Goroutine 1: Failed to acquire lease: context deadline exceeded
Goroutine 2: Failed to acquire lease: context deadline exceeded
Goroutine 8: Failed to acquire lease: context deadline exceeded
Goroutine 7: Failed to acquire lease: context deadline exceeded
Goroutine 6: Failed to acquire lease: context deadline exceeded
Goroutine 4: Failed to acquire lease: context deadline exceeded
Goroutine 3: Failed to acquire lease: context deadline exceeded
Test 3: Successfully limited concurrency to 3 (success: 3, errors: 7)
PASS
*/