95 lines
2.1 KiB
Go
95 lines
2.1 KiB
Go
|
package client
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"net/url"
|
||
|
|
"strconv"
|
||
|
|
"sync"
|
||
|
|
|
||
|
|
"github.com/jackc/puddle/v2"
|
||
|
|
"github.com/keegancsmith/rpc"
|
||
|
|
"github.com/pkg/errors"
|
||
|
|
"gitlab.com/wpetit/goweb/logger"
|
||
|
|
)
|
||
|
|
|
||
|
|
func NewClientPool(serverURL *url.URL, poolSize int) (*puddle.Pool[*rpc.Client], error) {
|
||
|
|
constructor := func(context.Context) (*rpc.Client, error) {
|
||
|
|
client, err := rpc.DialHTTPPath("tcp", serverURL.Host, serverURL.Path+"?"+serverURL.RawQuery)
|
||
|
|
if err != nil {
|
||
|
|
return nil, errors.WithStack(err)
|
||
|
|
}
|
||
|
|
|
||
|
|
return client, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
destructor := func(client *rpc.Client) {
|
||
|
|
if err := client.Close(); err != nil {
|
||
|
|
logger.Error(context.Background(), "could not close client", logger.CapturedE(errors.WithStack(err)))
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
maxPoolSize := int32(poolSize)
|
||
|
|
|
||
|
|
pool, err := puddle.NewPool(&puddle.Config[*rpc.Client]{Constructor: constructor, Destructor: destructor, MaxSize: maxPoolSize})
|
||
|
|
if err != nil {
|
||
|
|
return nil, errors.WithStack(err)
|
||
|
|
}
|
||
|
|
|
||
|
|
return pool, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
type WithClientFunc func(ctx context.Context, fn func(ctx context.Context, client *rpc.Client) error) error
|
||
|
|
|
||
|
|
func WithPooledClient(serverURL *url.URL) WithClientFunc {
|
||
|
|
var (
|
||
|
|
pool *puddle.Pool[*rpc.Client]
|
||
|
|
createPool sync.Once
|
||
|
|
)
|
||
|
|
|
||
|
|
return func(ctx context.Context, fn func(context.Context, *rpc.Client) error) error {
|
||
|
|
var err error
|
||
|
|
createPool.Do(func() {
|
||
|
|
rawPoolSize := serverURL.Query().Get("clientPoolSize")
|
||
|
|
if rawPoolSize == "" {
|
||
|
|
rawPoolSize = "5"
|
||
|
|
}
|
||
|
|
|
||
|
|
var poolSize int64
|
||
|
|
|
||
|
|
poolSize, err = strconv.ParseInt(rawPoolSize, 10, 32)
|
||
|
|
if err != nil {
|
||
|
|
err = errors.Wrap(err, "could not parse clientPoolSize url query parameter")
|
||
|
|
|
||
|
|
return
|
||
|
|
}
|
||
|
|
|
||
|
|
pool, err = NewClientPool(serverURL, int(poolSize))
|
||
|
|
if err != nil {
|
||
|
|
err = errors.WithStack(err)
|
||
|
|
|
||
|
|
return
|
||
|
|
}
|
||
|
|
})
|
||
|
|
if err != nil {
|
||
|
|
return errors.WithStack(err)
|
||
|
|
}
|
||
|
|
|
||
|
|
clientResource, err := pool.Acquire(ctx)
|
||
|
|
if err != nil {
|
||
|
|
return errors.WithStack(err)
|
||
|
|
}
|
||
|
|
|
||
|
|
if err := fn(ctx, clientResource.Value()); err != nil {
|
||
|
|
if errors.Is(err, rpc.ErrShutdown) {
|
||
|
|
clientResource.Destroy()
|
||
|
|
}
|
||
|
|
|
||
|
|
return errors.WithStack(err)
|
||
|
|
}
|
||
|
|
|
||
|
|
clientResource.Release()
|
||
|
|
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
}
|