gitee.com/h79/goutils@v1.22.10/thrift/client/pool.go (about)

     1  package client
     2  
     3  import (
     4  	"fmt"
     5  	"github.com/apache/thrift/lib/go/thrift"
     6  	"sync"
     7  )
     8  
     9  var (
    10  	pool *Pool
    11  )
    12  
    13  type Pool struct {
    14  	sync.RWMutex
    15  	client  map[string]*Client //有几个不同的机群,它们之间没有联系
    16  	factory map[string]thrift.TProtocolFactory
    17  }
    18  
    19  func newPool() *Pool {
    20  	if pool != nil {
    21  		return pool
    22  	}
    23  	pool = &Pool{
    24  		client:  make(map[string]*Client, 0),
    25  		factory: make(map[string]thrift.TProtocolFactory, 0),
    26  	}
    27  	return pool
    28  }
    29  
    30  func NewClient(cfg Config) (*Client, error) {
    31  	p := newPool()
    32  	return p.newClient(cfg)
    33  }
    34  
    35  func (p *Pool) newClient(cfg Config) (*Client, error) {
    36  	p.Lock()
    37  	defer p.Unlock()
    38  
    39  	f, err := p.createFactory(cfg.Protocol)
    40  	if err != nil {
    41  		return nil, err
    42  	}
    43  
    44  	return p.createClient(f, cfg.Target)
    45  }
    46  
    47  func (p *Pool) createFactory(protocol string) (thrift.TProtocolFactory, error) {
    48  
    49  	if f, ok := p.factory[protocol]; ok {
    50  		return f, nil
    51  	}
    52  
    53  	f, err := newFactory(protocol)
    54  	if err != nil {
    55  		return nil, err
    56  	}
    57  	p.factory[protocol] = f
    58  	return f, nil
    59  }
    60  
    61  func (p *Pool) createClient(tpf thrift.TProtocolFactory, target string) (*Client, error) {
    62  
    63  	if cli, ok := p.client[target]; ok {
    64  		return cli, nil
    65  	}
    66  	cli, err := newClient(tpf, target)
    67  	if err != nil {
    68  		return nil, err
    69  	}
    70  
    71  	p.client[target] = cli
    72  
    73  	return cli, nil
    74  }
    75  
    76  func newFactory(protocol string) (thrift.TProtocolFactory, error) {
    77  
    78  	var factory thrift.TProtocolFactory
    79  	var err error
    80  
    81  	switch protocol {
    82  	case "compact":
    83  		factory = thrift.NewTCompactProtocolFactoryConf(&thrift.TConfiguration{})
    84  		break
    85  	case "simple_json":
    86  		factory = thrift.NewTCompactProtocolFactoryConf(&thrift.TConfiguration{})
    87  		break
    88  	case "json":
    89  		factory = thrift.NewTJSONProtocolFactory()
    90  		break
    91  	case "binary", "":
    92  		factory = thrift.NewTBinaryProtocolFactoryConf(&thrift.TConfiguration{})
    93  		break
    94  	default:
    95  		err = fmt.Errorf("NOT support %v protocol factory", protocol)
    96  	}
    97  	return factory, err
    98  }