pool.go

package zdb

import (
	"context"
	"errors"
	"fmt"
	"github.com/jackc/pgx/v5/pgxpool"
	"log"
	"reflect"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

type Pool struct {
	ctx    context.Context
	logger Logger

	srvMaster      *conn
	srvSlaves      []*conn
	srvSlavesAsync []*conn

	mu *sync.RWMutex

	notAliveConns []*conn

	slavesIter      *atomic.Int64
	slavesAsyncIter *atomic.Int64
}

func NewDefault() *Pool {
	return &Pool{ctx: context.Background(), logger: log.Default()}
}

func (d *Pool) WithContext(ctx context.Context) *Pool {
	return &Pool{
		ctx:             ctx,
		srvMaster:       d.srvMaster,
		srvSlaves:       d.srvSlaves,
		srvSlavesAsync:  d.srvSlavesAsync,
		mu:              d.mu,
		notAliveConns:   d.notAliveConns,
		slavesIter:      d.slavesIter,
		slavesAsyncIter: d.slavesAsyncIter,
	}
}

func (d *Pool) NewConn(mode connMode, pgConnString string) error {
	if c, err := d.newConn(mode, pgConnString); err != nil {
		return err
	} else {
		switch mode {
		case ConnModeMaster:
			d.srvMaster = c
		case ConnModeSync:
			d.srvSlaves = append(d.srvSlaves, c)
		case ConnModeAsync:
			d.srvSlavesAsync = append(d.srvSlavesAsync, c)
		default:
			return errors.New("unknown mode")
		}

		return nil
	}
}
func (d *Pool) newConn(mode connMode, pgConnString string) (c *conn, err error) {
	var pool *pgxpool.Pool

	if !strings.Contains(pgConnString, "default_query_exec_mode=") {
		pgConnString += " default_query_exec_mode=simple_protocol"
	}

	if pool, err = pgxpool.New(d.ctx, pgConnString); err != nil {
		return &conn{Pool: pool, Alive: false, Mode: mode}, err
	}

	c = &conn{Pool: pool, Alive: false, Mode: mode}

	if err = d.testConn(c); err != nil {
		return c, err
	}

	c.Alive = true

	return c, nil
}

func (d *Pool) slave() *conn {
	if len(d.srvSlaves) == 0 {
		return d.srvMaster
	}

	d.mu.RLock()
	defer d.mu.RUnlock()

	return d.srvSlaves[d.slavesIter.Add(1)%int64(len(d.srvSlaves))]
}

func (d *Pool) slaveAsync() *conn {
	if len(d.srvSlavesAsync) == 0 {
		return d.slave()
	}

	d.mu.RLock()
	defer d.mu.RUnlock()

	return d.srvSlavesAsync[d.slavesAsyncIter.Add(1)%int64(len(d.srvSlavesAsync))]
}

func (d *Pool) execWrapper(pool connMode, dst any, f func(conn *conn, dst1 any) error) error {
	for {
		var c *conn

		if pool == ConnModeSync {
			c = d.slave()
		} else {
			c = d.slaveAsync()
		}

		if err := f(c, dst); err != nil {
			if c.Mode == ConnModeMaster {
				return err
			} else {
				if strings.Contains(err.Error(), "connect") || strings.Contains(err.Error(), "EOF") {
					d.logger.Printf("DB_EXEC_WRAPPER_ERR: %s", err.Error())
					d.setNotAliveConn(c)
					continue
				} else {
					return err
				}
			}
		}

		return nil
	}
}

func (d *Pool) testConn(q *conn) (err error) {
	var n any

	if err = d.qGet(q, &n, "SELECT 1"); err != nil {
		d.logger.Printf("DB_TEST_CONN_ERR: MODE: %v, HOST: %s, PORT: %d, ERR: %v",
			q.Mode.String(),
			q.Config().ConnConfig.Host,
			q.Config().ConnConfig.Port,
			err,
		)
	}

	return
}

func (d *Pool) setNotAliveConn(conn *conn) {
	d.mu.Lock()
	defer d.mu.Unlock()

	for i, slave := range d.srvSlaves {
		if slave == conn {
			conn.Alive = false
			d.notAliveConns = append(d.notAliveConns, conn)
			d.srvSlaves = remove(d.srvSlaves, i)
			return
		}
	}

	for i, slave := range d.srvSlavesAsync {
		if slave == conn {
			conn.Alive = false
			d.notAliveConns = append(d.notAliveConns, conn)
			d.srvSlavesAsync = remove(d.srvSlavesAsync, i)

			return
		}
	}
}

func (d *Pool) Test() {
	go func() {
		for {
			for i, p := range d.notAliveConns {
				if err := d.testConn(p); err == nil {
					d.mu.Lock()
					p.Alive = true
					d.notAliveConns = remove(d.notAliveConns, i)
					if p.Mode == ConnModeSync {
						d.srvSlaves = append(d.srvSlaves, p)
					} else if p.Mode == ConnModeAsync {
						d.srvSlavesAsync = append(d.srvSlavesAsync, p)
					}
					d.mu.Unlock()
				}
			}

			d.srvMaster.Alive = d.testConn(d.srvMaster) == nil

			time.Sleep(time.Second * 1)
		}
	}()
}

func (d *Pool) prepare(sql string, param map[string]any) string {
	for n, t := range param {
		switch t.(type) {
		case time.Time:
			sql = strings.ReplaceAll(sql, ":"+n+":", fmt.Sprintf("'%s'", t.(time.Time).UTC().Format(time.DateTime)))
		case *time.Time:
			if t.(*time.Time) == nil {
				sql = strings.ReplaceAll(sql, ":"+n+":", "NULL")
			} else {
				sql = strings.ReplaceAll(sql, ":"+n+":", fmt.Sprintf("'%s'", t.(*time.Time).UTC().Format(time.DateTime)))
			}
		case nil:
			sql = strings.ReplaceAll(sql, ":"+n+":", "NULL")
		case bool:
			if t.(bool) {
				sql = strings.ReplaceAll(sql, ":"+n+":", "TRUE")
			} else {
				sql = strings.ReplaceAll(sql, ":"+n+":", "FALSE")
			}
		case *string:
			if t.(*string) == nil || *t.(*string) == "NULL" {
				sql = strings.ReplaceAll(sql, ":"+n+":", "NULL")
			} else {
				sql = strings.ReplaceAll(sql, ":"+n+":", fmt.Sprintf("'%v'", strings.ReplaceAll(*t.(*string), "'", "''")))
			}
		case string:
			if t.(string) == "NULL" {
				sql = strings.ReplaceAll(sql, ":"+n+":", "NULL")
			} else {
				sql = strings.ReplaceAll(sql, ":"+n+":", fmt.Sprintf("'%v'", strings.ReplaceAll(t.(string), "'", "''")))
			}
		case *int:
			if t.(*int) == nil {
				sql = strings.ReplaceAll(sql, ":"+n+":", "NULL")
			} else {
				sql = strings.ReplaceAll(sql, ":"+n+":", fmt.Sprintf("%v", *t.(*int)))
			}
		case *bool:
			if t.(*bool) == nil {
				sql = strings.ReplaceAll(sql, ":"+n+":", "NULL")
			} else {
				if *t.(*bool) {
					sql = strings.ReplaceAll(sql, ":"+n+":", "TRUE")
				} else {
					sql = strings.ReplaceAll(sql, ":"+n+":", "FALSE")
				}
			}
		case *int64:
			if t.(*int64) == nil {
				sql = strings.ReplaceAll(sql, ":"+n+":", "NULL")
			} else {
				sql = strings.ReplaceAll(sql, ":"+n+":", fmt.Sprintf("%v", *t.(*int64)))
			}
		case *float64:
			if t.(*float64) == nil {
				sql = strings.ReplaceAll(sql, ":"+n+":", "NULL")
			} else {
				sql = strings.ReplaceAll(sql, ":"+n+":", fmt.Sprintf("%v", *t.(*float64)))
			}
		case int, int64:
			sql = strings.ReplaceAll(sql, ":"+n+":", fmt.Sprintf("%v", t))
		default:
			switch reflect.TypeOf(t).Kind() {
			case reflect.Slice:
				sql = strings.ReplaceAll(sql, ":"+n+":", "'{"+strings.Trim(strings.Join(strings.Split(fmt.Sprint(t), " "), ","), "[]")+"}'")
			}

			sql = strings.ReplaceAll(sql, ":"+n+":", fmt.Sprintf("'%v'", t))
		}
	}

	if _, ok := param["_debug"]; ok {
		d.logger.Printf(sql)
	}

	return sql
}