client_grpc.go

package boomer

import (
	"context"
	"fmt"
	"runtime"
	"sync"
	"sync/atomic"
	"time"

	"github.com/pkg/errors"
	"github.com/rs/zerolog/log"
	"golang.org/x/oauth2"
	"google.golang.org/grpc"
	"google.golang.org/grpc/backoff"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/credentials/oauth"
	"google.golang.org/grpc/metadata"

	"github.com/httprunner/boomer/data"
	"github.com/httprunner/boomer/grpc/messager"
)

type grpcClient struct {
	messager.MessageClient
	masterHost string
	masterPort int
	identity   string // nodeID

	config *grpcClientConfig

	fromMaster       chan *genericMessage
	toMaster         chan *genericMessage
	disconnectedChan chan bool
	shutdownChan     chan bool

	failCount int32
}

type grpcClientConfig struct {
	// ctx is used for the lifetime of the stream that may need to be canceled
	// on client shutdown.
	ctx       context.Context
	ctxCancel context.CancelFunc
	conn      *grpc.ClientConn
	biStream  messager.Message_BidirectionalStreamingMessageClient

	mutex sync.RWMutex
}

const token = "httprunner-secret-token"

// unaryInterceptor is an example unary interceptor.
func unaryInterceptor(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
	var credsConfigured bool
	for _, o := range opts {
		_, ok := o.(grpc.PerRPCCredsCallOption)
		if ok {
			credsConfigured = true
			break
		}
	}
	if !credsConfigured {
		opts = append(opts, grpc.PerRPCCredentials(oauth.NewOauthAccess(&oauth2.Token{
			AccessToken: token,
		})))
	}
	start := time.Now()
	err := invoker(ctx, method, req, reply, cc, opts...)
	end := time.Now()
	logger("RPC: %s, start time: %s, end time: %s, err: %v", method, start.Format("Basic"), end.Format(time.RFC3339), err)
	return err
}

// wrappedStream  wraps around the embedded grpc.ClientStream, and intercepts the RecvMsg and
// SendMsg method call.
type wrappedStream struct {
	grpc.ClientStream
}

func (w *wrappedStream) RecvMsg(m interface{}) error {
	logger("Receive a message (Type: %T) at %v", m, time.Now().Format(time.RFC3339))
	return w.ClientStream.RecvMsg(m)
}

func (w *wrappedStream) SendMsg(m interface{}) error {
	logger("Send a message (Type: %T) at %v", m, time.Now().Format(time.RFC3339))
	return w.ClientStream.SendMsg(m)
}

func newWrappedStream(s grpc.ClientStream) grpc.ClientStream {
	return &wrappedStream{s}
}

func extractToken(ctx context.Context) (tkn string, ok bool) {
	md, ok := metadata.FromIncomingContext(ctx)
	if !ok || len(md[token]) == 0 {
		return "", false
	}

	return md[token][0], true
}

// streamInterceptor is an example stream interceptor.
func streamInterceptor(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
	var credsConfigured bool
	for _, o := range opts {
		_, ok := o.(*grpc.PerRPCCredsCallOption)
		if ok {
			credsConfigured = true
			break
		}
	}
	if !credsConfigured {
		opts = append(opts, grpc.PerRPCCredentials(oauth.NewOauthAccess(&oauth2.Token{
			AccessToken: token,
		})))
	}
	s, err := streamer(ctx, desc, cc, method, opts...)
	if err != nil {
		return nil, err
	}
	return newWrappedStream(s), nil
}

func (c *grpcClientConfig) getBiStreamClient() messager.Message_BidirectionalStreamingMessageClient {
	c.mutex.RLock()
	defer c.mutex.RUnlock()
	return c.biStream
}

func (c *grpcClientConfig) setBiStreamClient(s messager.Message_BidirectionalStreamingMessageClient) {
	c.mutex.Lock()
	defer c.mutex.Unlock()
	c.biStream = s
}

func newClient(masterHost string, masterPort int, identity string) (client *grpcClient) {
	log.Info().Msg("Boomer is built with grpc support.")
	// Initiate the stream with a context that supports cancellation.
	ctx, cancel := context.WithCancel(context.Background())
	client = &grpcClient{
		masterHost:       masterHost,
		masterPort:       masterPort,
		identity:         identity,
		fromMaster:       make(chan *genericMessage, 100),
		toMaster:         make(chan *genericMessage, 100),
		disconnectedChan: make(chan bool),
		shutdownChan:     make(chan bool),
		config: &grpcClientConfig{
			ctx:       ctx,
			ctxCancel: cancel,
			mutex:     sync.RWMutex{},
		},
	}
	return client
}

func (c *grpcClient) start() (err error) {
	addr := fmt.Sprintf("%v:%v", c.masterHost, c.masterPort)
	// Create tls based credential.
	creds, err := credentials.NewClientTLSFromFile(data.Path("x509/ca_cert.pem"), "www.httprunner.com")
	if err != nil {
		log.Fatal().Msg(fmt.Sprintf("failed to load credentials: %v", err))
	}
	opts := []grpc.DialOption{
		// oauth.NewOauthAccess requires the configuration of transport
		// credentials.
		grpc.WithTransportCredentials(creds),
		grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(32 * 10e9)),
		grpc.WithUnaryInterceptor(unaryInterceptor),
		grpc.WithStreamInterceptor(streamInterceptor),
		grpc.WithConnectParams(grpc.ConnectParams{
			Backoff: backoff.Config{
				BaseDelay:  1 * time.Second,
				Multiplier: 1.2,
				MaxDelay:   3 * time.Second,
			},
			MinConnectTimeout: 3 * time.Second,
		}),
	}
	c.config.conn, err = grpc.Dial(addr, opts...)
	if err != nil {
		log.Error().Err(err).Msg("failed to connect")
		return err
	}
	c.MessageClient = messager.NewMessageClient(c.config.conn)
	return nil
}

func (c *grpcClient) register(ctx context.Context) error {
	ctx, cancel := context.WithTimeout(ctx, time.Second)
	defer cancel()
	res, err := c.Register(ctx, &messager.RegisterRequest{NodeID: c.identity, Os: runtime.GOOS, Arch: runtime.GOARCH})
	if err != nil {
		return err
	}
	if res.Code != "0" {
		return errors.New(res.Message)
	}
	return nil
}

func (c *grpcClient) signOut(ctx context.Context) error {
	ctx, cancel := context.WithTimeout(ctx, time.Second)
	defer cancel()

	res, err := c.SignOut(ctx, &messager.SignOutRequest{NodeID: c.identity})
	if err != nil {
		return err
	}
	if res.Code != "0" {
		return errors.New(res.Message)
	}
	return nil
}

func (c *grpcClient) newBiStreamClient() (err error) {
	md := metadata.New(map[string]string{token: c.identity})
	ctx := metadata.NewOutgoingContext(c.config.ctx, md)
	biStream, err := c.BidirectionalStreamingMessage(ctx)
	if err != nil {
		return err
	}
	// reset failCount
	atomic.StoreInt32(&c.failCount, 0)
	// set bidirectional stream client
	c.config.setBiStreamClient(biStream)
	println("successful to establish bidirectional stream with master, press Ctrl+c to quit.")
	return nil
}

func (c *grpcClient) recvChannel() chan *genericMessage {
	return c.fromMaster
}

func (c *grpcClient) recv() {
	for {
		select {
		case <-c.shutdownChan:
			return
		default:
			if c.config.getBiStreamClient() == nil {
				time.Sleep(1 * time.Second)
				continue
			}
			msg, err := c.config.getBiStreamClient().Recv()
			if err != nil {
				time.Sleep(1 * time.Second)
				// log.Error().Err(err).Msg("failed to get message")
				continue
			}
			if msg == nil {
				continue
			}

			if msg.NodeID != c.identity {
				log.Info().
					Str("nodeID", msg.NodeID).
					Str("type", msg.Type).
					Interface("data", msg.Data).
					Msg(fmt.Sprintf("not for me(%s)", c.identity))
				continue
			}

			c.fromMaster <- &genericMessage{
				Type:    msg.Type,
				Profile: msg.Profile,
				Data:    msg.Data,
				NodeID:  msg.NodeID,
				Tasks:   msg.Tasks,
			}

			log.Info().
				Str("nodeID", msg.NodeID).
				Str("type", msg.Type).
				Interface("data", msg.Data).
				Interface("tasks", msg.Tasks).
				Msg("receive data from master")
		}
	}
}

func (c *grpcClient) sendChannel() chan *genericMessage {
	return c.toMaster
}

func (c *grpcClient) send() {
	for {
		select {
		case <-c.shutdownChan:
			return
		case msg := <-c.toMaster:
			c.sendMessage(msg)

			// We may send genericMessage to master.
			switch msg.Type {
			case "quit":
				c.disconnectedChan <- true
			}
		}
	}
}

func (c *grpcClient) sendMessage(msg *genericMessage) {
	log.Info().
		Str("nodeID", msg.NodeID).
		Str("type", msg.Type).
		Interface("data", msg.Data).
		Msg("send data to server")
	if c.config.getBiStreamClient() == nil {
		atomic.AddInt32(&c.failCount, 1)
		return
	}
	err := c.config.getBiStreamClient().Send(&messager.StreamRequest{Type: msg.Type, Data: msg.Data, NodeID: msg.NodeID})
	if err == nil {
		atomic.StoreInt32(&c.failCount, 0)
		return
	}
	// log.Error().Err(err).Interface("genericMessage", *msg).Msg("failed to send message")
	if msg.Type == "heartbeat" {
		atomic.AddInt32(&c.failCount, 1)
	}
}

func (c *grpcClient) disconnectedChannel() chan bool {
	return c.disconnectedChan
}

func (c *grpcClient) close() {
	close(c.shutdownChan)
	c.config.ctxCancel()
	if c.config.conn != nil {
		c.config.conn.Close()
	}
}