pool.go

package gomail

import (
	"context"
	"errors"
	"sync"
	"time"
)

var (
	ErrPoolClosed = errors.New("connection pool is closed")
	ErrPoolEmpty  = errors.New("no connections available")
)

// Pool manages a pool of persistent SMTP connections
type Pool struct {
	gomail       *GoMail
	connections  chan *SendCloser
	mu           sync.Mutex
	closed       bool
	maxSize      int
	currentSize  int
	idleTimeout  time.Duration
	maxLifetime  time.Duration
}

// PoolConfig configures the connection pool
type PoolConfig struct {
	MaxSize     int           // Maximum number of connections in the pool
	IdleTimeout time.Duration // How long a connection can be idle before being closed
	MaxLifetime time.Duration // Maximum lifetime of a connection
}

// NewPool creates a new connection pool
func (g *GoMail) NewPool(config PoolConfig) *Pool {
	if config.MaxSize <= 0 {
		config.MaxSize = 10
	}
	if config.IdleTimeout <= 0 {
		config.IdleTimeout = 5 * time.Minute
	}
	if config.MaxLifetime <= 0 {
		config.MaxLifetime = 30 * time.Minute
	}

	return &Pool{
		gomail:      g,
		connections: make(chan *SendCloser, config.MaxSize),
		maxSize:     config.MaxSize,
		idleTimeout: config.IdleTimeout,
		maxLifetime: config.MaxLifetime,
	}
}

// Get retrieves a connection from the pool or creates a new one
func (p *Pool) Get(ctx context.Context) (*SendCloser, error) {
	p.mu.Lock()
	if p.closed {
		p.mu.Unlock()
		return nil, ErrPoolClosed
	}
	p.mu.Unlock()

	select {
	case conn := <-p.connections:
		// Got a connection from the pool
		return conn, nil
	default:
		// No available connections, try to create a new one
		p.mu.Lock()
		if p.currentSize < p.maxSize {
			p.currentSize++
			p.mu.Unlock()
			return p.gomail.Dial()
		}
		p.mu.Unlock()

		// Wait for a connection to become available
		select {
		case conn := <-p.connections:
			return conn, nil
		case <-ctx.Done():
			return nil, ctx.Err()
		}
	}
}

// Put returns a connection to the pool
func (p *Pool) Put(conn *SendCloser) error {
	p.mu.Lock()
	defer p.mu.Unlock()

	if p.closed {
		conn.Close()
		return ErrPoolClosed
	}

	select {
	case p.connections <- conn:
		return nil
	default:
		// Pool is full, close the connection
		conn.Close()
		p.currentSize--
		return nil
	}
}

// Send sends an email using a connection from the pool
func (p *Pool) Send(ctx context.Context, mail *Mail) error {
	conn, err := p.Get(ctx)
	if err != nil {
		return err
	}

	sendErr := conn.Send(ctx, mail)

	// Return connection to pool regardless of send result
	putErr := p.Put(conn)

	if sendErr != nil {
		return sendErr
	}
	return putErr
}

// SendBulk sends multiple emails using a connection from the pool
func (p *Pool) SendBulk(ctx context.Context, mails []*Mail) error {
	if len(mails) == 0 {
		return nil
	}

	conn, err := p.Get(ctx)
	if err != nil {
		return err
	}

	var sendErr error
	for _, mail := range mails {
		if err := conn.Send(ctx, mail); err != nil {
			sendErr = err
			break
		}
	}

	// Return connection to pool
	putErr := p.Put(conn)

	if sendErr != nil {
		return sendErr
	}
	return putErr
}

// Close closes all connections in the pool
func (p *Pool) Close() error {
	p.mu.Lock()
	defer p.mu.Unlock()

	if p.closed {
		return nil
	}

	p.closed = true
	close(p.connections)

	// Close all connections in the pool
	for conn := range p.connections {
		conn.Close()
	}

	return nil
}

// Size returns the current number of connections in the pool
func (p *Pool) Size() int {
	p.mu.Lock()
	defer p.mu.Unlock()
	return p.currentSize
}