edge/pkg/storage/driver/rpc/client/client_pool.go

134 lines
2.9 KiB
Go

package client
import (
"context"
"net/url"
"strconv"
"sync"
"time"
"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)
}
attempts := 0
max := 5
for {
if attempts >= max {
logger.Debug(ctx, "rpc client call retrying failed", logger.F("attempts", attempts))
return errors.Wrapf(err, "rpc client call failed after %d attempts", max)
}
clientResource, err := pool.Acquire(ctx)
if err != nil {
return errors.WithStack(err)
}
client := clientResource.Value()
if err := fn(ctx, client); err != nil {
if errors.Is(err, rpc.ErrShutdown) {
clientResource.Destroy()
wait := time.Duration(8<<(attempts+1)) * time.Millisecond
logger.Warn(
ctx, "rpc client connection is shutdown, retrying",
logger.F("attempts", attempts),
logger.F("max", max),
logger.F("delay", wait),
)
timer := time.NewTimer(wait)
select {
case <-timer.C:
attempts++
continue
case <-ctx.Done():
if err := ctx.Err(); err != nil {
return errors.WithStack(err)
}
return nil
}
}
clientResource.Release()
return errors.WithStack(err)
}
clientResource.Release()
return nil
}
}
}