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

     1  package client
     2  
     3  import (
     4  	"fmt"
     5  	"gitee.com/h79/goutils/common/server"
     6  	"gitee.com/h79/goutils/common/system"
     7  	"gitee.com/h79/goutils/thrift/resolver"
     8  	"gitee.com/h79/goutils/thrift/thriftutil"
     9  	"github.com/apache/thrift/lib/go/thrift"
    10  	"sync"
    11  )
    12  
    13  var TransInvalidateERR = fmt.Errorf("transport invalidate")
    14  
    15  type Client struct {
    16  	target       string
    17  	parsedTarget resolver.Target
    18  
    19  	firstResolveEvent *system.Event
    20  	mu                sync.RWMutex
    21  	resolverWrapper   *ccResolverWrapper
    22  
    23  	tpf        thrift.TProtocolFactory
    24  	tTransport thrift.TTransport
    25  	tClient    thrift.TClient
    26  }
    27  
    28  func newClient(tpf thrift.TProtocolFactory, target string) (*Client, error) {
    29  	cli := &Client{
    30  		tpf:    tpf,
    31  		target: target,
    32  	}
    33  	if err := cli.newTransport(); err != nil {
    34  		return nil, err
    35  	}
    36  	return cli, nil
    37  }
    38  
    39  func (c *Client) Open() error {
    40  	if c.tTransport != nil {
    41  		return c.tTransport.Open()
    42  	}
    43  	return c.build()
    44  }
    45  
    46  func (c *Client) Close() error {
    47  	if c.tTransport != nil {
    48  		return c.tTransport.Close()
    49  	}
    50  	return TransInvalidateERR
    51  }
    52  
    53  func (c *Client) IsOpen() bool {
    54  	if c.tTransport != nil {
    55  		return c.tTransport.IsOpen()
    56  	}
    57  	return false
    58  }
    59  
    60  func (c *Client) FramedTransport() {
    61  	if c.tTransport != nil {
    62  		c.tTransport = thrift.NewTFramedTransport(c.tTransport)
    63  	}
    64  }
    65  
    66  func (c *Client) SetHeader(headers map[string]string) {
    67  	if len(headers) == 0 {
    68  		panic(fmt.Errorf(""))
    69  		return
    70  	}
    71  	if httpTrans, ok := c.tTransport.(*thrift.THttpClient); ok {
    72  		for key, value := range headers {
    73  			httpTrans.SetHeader(key, value)
    74  		}
    75  	}
    76  }
    77  
    78  func (c *Client) In() thrift.TProtocol {
    79  	if c.tTransport != nil {
    80  		return nil
    81  	}
    82  	return c.tpf.GetProtocol(c.tTransport)
    83  }
    84  
    85  func (c *Client) Out() thrift.TProtocol {
    86  	return c.In()
    87  }
    88  
    89  func (c *Client) TClient() thrift.TClient {
    90  	if c.tTransport != nil {
    91  		return thrift.NewTStandardClient(c.In(), c.Out())
    92  	}
    93  	return nil
    94  }
    95  
    96  func (c *Client) getResolver(scheme string) resolver.Builder {
    97  	if scheme == "" {
    98  		return nil
    99  	}
   100  	return resolver.Get(scheme)
   101  }
   102  
   103  func (c *Client) build() error {
   104  	return c.newTransport()
   105  }
   106  
   107  func (c *Client) newTransport() error {
   108  
   109  	c.parsedTarget = thriftutil.ParseTarget(c.target)
   110  	builder := c.getResolver(c.parsedTarget.Scheme)
   111  	if builder == nil {
   112  		addr, er := server.Parse(c.target)
   113  		if er != nil {
   114  			return er
   115  		}
   116  		var err error
   117  		var trans thrift.TTransport
   118  
   119  		if addr.IsHttp() {
   120  			trans, err = thrift.NewTHttpClient(addr.To())
   121  		} else if (addr.HasIP() || addr.HasDomain()) && addr.HasPort() {
   122  			trans, err = thrift.NewTSocket(addr.To())
   123  		}
   124  		if err != nil {
   125  			return err
   126  		}
   127  		c.tTransport = trans
   128  	} else if c.resolverWrapper == nil {
   129  		resolverWrapper, err := newCCResolverWrapper(c, builder)
   130  		if err != nil {
   131  			return err
   132  		}
   133  		c.mu.Lock()
   134  		c.resolverWrapper = resolverWrapper
   135  		c.mu.Unlock()
   136  	} else {
   137  		c.mu.Lock()
   138  		r := c.resolverWrapper
   139  		c.mu.Unlock()
   140  		r.resolveNow()
   141  	}
   142  	return nil
   143  }
   144  
   145  func (c *Client) updateResolverState(s resolver.State, err error) error {
   146  	return nil
   147  }