github.com/1aal/kubeblocks@v0.0.0-20231107070852-e1c03e598921/pkg/lorry/grpcserver/gprc_server.go (about)

     1  /*
     2  Copyright (C) 2022-2023 ApeCloud Co., Ltd
     3  
     4  This file is part of KubeBlocks project
     5  
     6  This program is free software: you can redistribute it and/or modify
     7  it under the terms of the GNU Affero General Public License as published by
     8  the Free Software Foundation, either version 3 of the License, or
     9  (at your option) any later version.
    10  
    11  This program is distributed in the hope that it will be useful
    12  but WITHOUT ANY WARRANTY; without even the implied warranty of
    13  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    14  GNU Affero General Public License for more details.
    15  
    16  You should have received a copy of the GNU Affero General Public License
    17  along with this program.  If not, see <http://www.gnu.org/licenses/>.
    18  */
    19  
    20  package grpcserver
    21  
    22  import (
    23  	"context"
    24  	"encoding/json"
    25  	"flag"
    26  	"fmt"
    27  	"net"
    28  	"strings"
    29  
    30  	"github.com/go-logr/logr"
    31  	"github.com/pkg/errors"
    32  	"google.golang.org/grpc"
    33  	"google.golang.org/grpc/codes"
    34  	health "google.golang.org/grpc/health/grpc_health_v1"
    35  	"google.golang.org/grpc/status"
    36  	ctrl "sigs.k8s.io/controller-runtime"
    37  
    38  	"github.com/1aal/kubeblocks/pkg/lorry/operations"
    39  	"github.com/1aal/kubeblocks/pkg/lorry/util"
    40  )
    41  
    42  type GRPCServer struct {
    43  	logger             logr.Logger
    44  	checkRoleOperation operations.Operation
    45  }
    46  
    47  var (
    48  	grpcPort int
    49  )
    50  
    51  const (
    52  	DefaultGRPCPort = 50001
    53  )
    54  
    55  func init() {
    56  	flag.IntVar(&grpcPort, "grpcport", DefaultGRPCPort, "lorry grpc default port")
    57  }
    58  
    59  func (s *GRPCServer) Check(ctx context.Context, in *health.HealthCheckRequest) (*health.HealthCheckResponse, error) {
    60  	resp, err := s.checkRoleOperation.Do(ctx, nil)
    61  
    62  	var status = health.HealthCheckResponse_SERVING
    63  	if err != nil {
    64  		status = health.HealthCheckResponse_NOT_SERVING
    65  		if _, ok := err.(util.ProbeError); !ok {
    66  			s.logger.Error(err, "role probe failed")
    67  			return &health.HealthCheckResponse{Status: status}, err
    68  		} else {
    69  			body, _ := json.Marshal(resp.Data)
    70  			s.logger.Info("Role changed event detected", "role", string(body))
    71  			return &health.HealthCheckResponse{Status: status}, errors.New(string(body))
    72  		}
    73  	}
    74  
    75  	s.logger.Info("No event detected", "response", resp)
    76  	return &health.HealthCheckResponse{Status: status}, nil
    77  }
    78  
    79  func (s *GRPCServer) Watch(in *health.HealthCheckRequest, _ health.Health_WatchServer) error {
    80  	// didn't implement the `watch` function
    81  	return status.Error(codes.Unimplemented, "unimplemented")
    82  }
    83  
    84  func (s *GRPCServer) StartNonBlocking() error {
    85  	listen, err := net.Listen("tcp", fmt.Sprintf(":%d", grpcPort))
    86  	if err != nil {
    87  		return errors.Wrap(err, "grpc server listen failed")
    88  	}
    89  	server := grpc.NewServer()
    90  	health.RegisterHealthServer(server, s)
    91  
    92  	go func() {
    93  		err = server.Serve(listen)
    94  		if err != nil {
    95  			s.logger.Error(err, "grpcserver serve failed")
    96  		}
    97  	}()
    98  	return nil
    99  }
   100  
   101  func NewGRPCServer() (*GRPCServer, error) {
   102  	checkRoleOperation, ok := operations.Operations()[strings.ToLower(string(util.CheckRoleOperation))]
   103  	if !ok {
   104  		return nil, errors.New("check role operation not found")
   105  	}
   106  	err := checkRoleOperation.Init(context.Background())
   107  	if err != nil {
   108  		return nil, err
   109  	}
   110  
   111  	return &GRPCServer{
   112  		logger:             ctrl.Log.WithName("grpc"),
   113  		checkRoleOperation: checkRoleOperation,
   114  	}, nil
   115  }