gothriftpool.go

package gothriftpool

import (
	"bytes"
	"fmt"
	"go/format"
	"text/template"

	"github.com/koofr/gointerfacer"
)

const proxyTemplate = `// Autogenerated by gothriftpool generator
// DO NOT EDIT UNLESS YOU ARE SURE THAT YOU KNOW WHAT YOU ARE DOING

package {{.ProxyPackage}}

import (
	"context"

	"github.com/koofr/go-resourcepool"

	"{{.InterfaceImport}}"
)

type ErrCreateResource struct {
	Err error
}

func NewErrCreateResource(err error) *ErrCreateResource {
	return &ErrCreateResource{Err: err}
}

func (e *ErrCreateResource) Error() string {
	return "{{.ClientType}} resource could not be created: " + e.Err.Error()
}

func (e *ErrCreateResource) Unwrap() error {
	return e.Err
}

type ClientFactory func(context.Context) (client {{.ClientType}}, close func(), err error)

type resource struct {
	client {{.ClientType}}
	close func()
	isClosed bool
}

func (r *resource) Close() {
	r.isClosed = true
	r.close()
}

func (r *resource) IsClosed() bool {
	return r.isClosed
}

type ClientPool struct {
	clientFactory ClientFactory
	pool          *resourcepool.ResourcePool
	retries       int
}

func NewClientPool(clientFactory ClientFactory, idleCapacity int, maxResources int) *ClientPool {
	p := &ClientPool{
		clientFactory: clientFactory,
		retries:       2,
	}

	p.pool = resourcepool.NewResourcePool(p.createResource, idleCapacity, maxResources)

	return p
}

func (p *ClientPool) Close() {
	p.pool.Close()
}

func (p *ClientPool) SetRetries(retries int) {
	p.retries = retries
}

func (p *ClientPool) GetPool() *resourcepool.ResourcePool {
	return p.pool
}

func (p *ClientPool) createResource(ctx context.Context) (r resourcepool.Resource, err error) {
	client, close, err := p.clientFactory(ctx)
	if err != nil {
		return nil, err
	}

	r = &resource{
		client:   client,
		close:    close,
		isClosed: false,
	}

	return
}

func (p *ClientPool) GetClient(ctx context.Context) (client {{.ClientType}}, release func(), err error) {
	var lastErr error

	for i := 0; i < p.retries; i++ {
		r, err := p.pool.Acquire(ctx)
		if err != nil {
			if err == context.Canceled || err == context.DeadlineExceeded {
				return nil, nil, err
			}

			p.pool.Empty()

			lastErr = err
			continue
		}

		err = r.(*resource).client.Ping(ctx)
		if err != nil {
			r.Close()

			p.pool.Release(r)

			p.pool.Empty()

			lastErr = err
			continue
		}

		release = func() {
			if ctx.Err() != nil {
				r.Close()
			}

			p.pool.Release(r)
		}

		return r.(*resource).client, release, nil
	}

	return nil, nil, NewErrCreateResource(lastErr)
}

type Proxy struct {
	clientFactory ClientFactory
}

func New(clientFactory ClientFactory) *Proxy {
	return &Proxy{
		clientFactory: clientFactory,
	}
}

{{range .Functions}}
func (p *Proxy) {{.Name}}({{range .Params}}{{.Name}} {{.Type}}, {{end}})({{range .Res}}{{.Name}} {{.Type}}, {{end}}){
	thriftClient, thriftClientClose, err := p.clientFactory(ctx)
	if err != nil {
		return
	}
	defer thriftClientClose()

	return thriftClient.{{.Name}}({{range .Params}}{{.Name}}, {{end}})
}
{{end}}
`

type Generator struct {
	InterfaceImport  string
	InterfacePackage string
	InterfaceName    string
	ProxyPackage     string
	Functions        []gointerfacer.Func
}

type ThriftInterface struct {
	InterfaceImport string
	ProxyPackage    string
	ClientType      string
	Functions       []gointerfacer.Func
}

func NewGenerator(iface string) (g *Generator, err error) {
	ifaceImport, ifacePkg, ifaceName, err := gointerfacer.FindInterface(iface)

	if err != nil {
		return nil, err
	}

	proxyPackage := ifacePkg + "proxy"

	functions, err := gointerfacer.Functions(iface)

	if err != nil {
		return nil, err
	}

	var pingFunc *gointerfacer.Func

	for _, f := range functions {
		if f.Name == "Ping" &&
			len(f.Params) == 1 &&
			f.Params[0].Name == "ctx" &&
			f.Params[0].Type == "context.Context" &&
			len(f.Res) == 1 &&
			f.Res[0].Type == "error" {
			pingFunc = &f
		}
	}

	if pingFunc == nil {
		return nil, fmt.Errorf("Interface %s does not contain function Ping() error", iface)
	}

	g = &Generator{
		InterfaceImport:  ifaceImport,
		InterfacePackage: ifacePkg,
		InterfaceName:    ifaceName,
		ProxyPackage:     proxyPackage,
		Functions:        functions,
	}

	return g, nil
}

func (g *Generator) render() (source []byte, err error) {
	ti := &ThriftInterface{
		InterfaceImport: g.InterfaceImport,
		ProxyPackage:    g.ProxyPackage,
		ClientType:      g.InterfacePackage + "." + g.InterfaceName,
		Functions:       g.Functions,
	}

	tmpl, err := template.New("tpl").Parse(proxyTemplate)

	if err != nil {
		return nil, err
	}

	buf := &bytes.Buffer{}

	err = tmpl.Execute(buf, ti)

	if err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

func (g *Generator) format(source []byte) (code []byte, err error) {
	return format.Source(source)
}

func (g *Generator) Generate() (code []byte, err error) {
	source, err := g.render()

	if err != nil {
		return nil, err
	}

	code, err = g.format(source)

	if err != nil {
		return nil, err
	}

	return code, nil
}