traefik/vendor/github.com/ExpediaDotCom/haystack-client-go/remote_client.go

166 lines
4.1 KiB
Go
Raw Normal View History

2019-05-08 22:14:04 +00:00
/*
* Copyright 2018 Expedia Group.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*
*/
package haystack
import (
"bytes"
"fmt"
"io/ioutil"
"net/http"
"time"
"github.com/golang/protobuf/proto"
context "golang.org/x/net/context"
grpc "google.golang.org/grpc"
)
/*RemoteClient remote client*/
type RemoteClient interface {
Send(span *Span)
Close() error
SetLogger(logger Logger)
}
/*GrpcClient grpc client*/
type GrpcClient struct {
conn *grpc.ClientConn
client SpanAgentClient
timeout time.Duration
logger Logger
}
/*NewGrpcClient returns a new grpc client*/
func NewGrpcClient(host string, port int, timeout time.Duration) *GrpcClient {
targetHost := fmt.Sprintf("%s:%d", host, port)
conn, err := grpc.Dial(targetHost, grpc.WithInsecure())
if err != nil {
panic(fmt.Sprintf("fail to connect to agent with error: %v", err))
}
return &GrpcClient{
conn: conn,
client: NewSpanAgentClient(conn),
timeout: timeout,
}
}
/*Send a proto span to grpc server*/
func (c *GrpcClient) Send(span *Span) {
ctx, cancel := context.WithTimeout(context.Background(), c.timeout)
defer cancel()
result, err := c.client.Dispatch(ctx, span)
if err != nil {
c.logger.Error("Fail to dispatch to haystack-agent with error %v", err)
} else if result.GetCode() != DispatchResult_SUCCESS {
c.logger.Error(fmt.Sprintf("Fail to dispatch to haystack-agent with error code: %d, message :%s", result.GetCode(), result.GetErrorMessage()))
} else {
c.logger.Debug(fmt.Sprintf("span [%v] has been successfully dispatched to haystack", span))
}
}
/*Close the grpc client*/
func (c *GrpcClient) Close() error {
return c.conn.Close()
}
/*SetLogger sets the logger*/
func (c *GrpcClient) SetLogger(logger Logger) {
c.logger = logger
}
/*HTTPClient a http client*/
type HTTPClient struct {
url string
headers map[string]string
client *http.Client
logger Logger
}
/*NewHTTPClient returns a new http client*/
func NewHTTPClient(url string, headers map[string]string, timeout time.Duration) *HTTPClient {
httpClient := &http.Client{
Timeout: timeout,
}
return &HTTPClient{
url: url,
headers: headers,
client: httpClient,
}
}
/*Send a proto span to http server*/
func (c *HTTPClient) Send(span *Span) {
serializedBytes, marshalErr := proto.Marshal(span)
if marshalErr != nil {
c.logger.Error("Fail to serialize the span to proto bytes, error=%v", marshalErr)
return
}
postRequest, requestErr := http.NewRequest(http.MethodPost, c.url, bytes.NewReader(serializedBytes))
if requestErr != nil {
c.logger.Error("Fail to create request for posting span to haystack server, error=%v", requestErr)
return
}
if c.headers != nil {
for k, v := range c.headers {
postRequest.Header.Add(k, v)
}
}
resp, err := c.client.Do(postRequest)
if err != nil {
c.logger.Error("Fail to dispatch to haystack http server, error=%v", err)
}
defer func() {
closeErr := resp.Body.Close()
if closeErr != nil {
/* do nothing */
}
}()
respBytes, respErr := ioutil.ReadAll(resp.Body)
if respErr != nil {
c.logger.Error("Fail to read the http response from haystack server, error=%v", respErr)
return
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
c.logger.Error("Fail to dispatch the span to haystack http server with statusCode=%d , payload=%s", resp.StatusCode, string(respBytes))
} else {
c.logger.Debug(fmt.Sprintf("span [%v] has been successfully dispatched to haystack, response=%s", span, string(respBytes)))
}
}
/*Close the http client*/
func (c *HTTPClient) Close() error {
return nil
}
/*SetLogger sets the logger*/
func (c *HTTPClient) SetLogger(logger Logger) {
c.logger = logger
}