github.com/MetalBlockchain/metalgo@v1.11.9/vms/rpcchainvm/vm.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package rpcchainvm
     5  
     6  import (
     7  	"context"
     8  	"fmt"
     9  	"os"
    10  	"os/signal"
    11  	"syscall"
    12  	"time"
    13  
    14  	"google.golang.org/grpc"
    15  	"google.golang.org/grpc/health"
    16  
    17  	"github.com/MetalBlockchain/metalgo/snow/engine/snowman/block"
    18  	"github.com/MetalBlockchain/metalgo/utils"
    19  	"github.com/MetalBlockchain/metalgo/version"
    20  	"github.com/MetalBlockchain/metalgo/vms/rpcchainvm/grpcutils"
    21  	"github.com/MetalBlockchain/metalgo/vms/rpcchainvm/gruntime"
    22  	"github.com/MetalBlockchain/metalgo/vms/rpcchainvm/runtime"
    23  
    24  	vmpb "github.com/MetalBlockchain/metalgo/proto/pb/vm"
    25  	runtimepb "github.com/MetalBlockchain/metalgo/proto/pb/vm/runtime"
    26  	healthpb "google.golang.org/grpc/health/grpc_health_v1"
    27  )
    28  
    29  const defaultRuntimeDialTimeout = 5 * time.Second
    30  
    31  // The address of the Runtime server is expected to be passed via ENV `runtime.EngineAddressKey`.
    32  // This address is used by the Runtime client to send Initialize RPC to server.
    33  //
    34  // Serve starts the RPC Chain VM server and performs a handshake with the VM runtime service.
    35  func Serve(ctx context.Context, vm block.ChainVM, opts ...grpcutils.ServerOption) error {
    36  	signals := make(chan os.Signal, 2)
    37  	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)
    38  	defer signal.Stop(signals)
    39  
    40  	var allowShutdown utils.Atomic[bool]
    41  	server := newVMServer(vm, &allowShutdown, opts...)
    42  	go func(ctx context.Context) {
    43  		defer func() {
    44  			server.GracefulStop()
    45  			fmt.Println("vm server: graceful termination success")
    46  		}()
    47  
    48  		for {
    49  			select {
    50  			case s := <-signals:
    51  				// We drop all signals until our parent process has notified us
    52  				// that we are shutting down. Once we are in the shutdown
    53  				// workflow, we will gracefully exit upon receiving a SIGTERM.
    54  				if !allowShutdown.Get() {
    55  					fmt.Printf("runtime engine: ignoring signal: %s\n", s)
    56  					continue
    57  				}
    58  
    59  				switch s {
    60  				case syscall.SIGINT:
    61  					fmt.Printf("runtime engine: ignoring signal: %s\n", s)
    62  				case syscall.SIGTERM:
    63  					fmt.Printf("runtime engine: received shutdown signal: %s\n", s)
    64  					return
    65  				}
    66  			case <-ctx.Done():
    67  				fmt.Println("runtime engine: context has been cancelled")
    68  				return
    69  			}
    70  		}
    71  	}(ctx)
    72  
    73  	// address of Runtime server from ENV
    74  	runtimeAddr := os.Getenv(runtime.EngineAddressKey)
    75  	if runtimeAddr == "" {
    76  		return fmt.Errorf("required env var missing: %q", runtime.EngineAddressKey)
    77  	}
    78  
    79  	clientConn, err := grpcutils.Dial(runtimeAddr)
    80  	if err != nil {
    81  		return fmt.Errorf("failed to create client conn: %w", err)
    82  	}
    83  
    84  	client := gruntime.NewClient(runtimepb.NewRuntimeClient(clientConn))
    85  
    86  	listener, err := grpcutils.NewListener()
    87  	if err != nil {
    88  		return fmt.Errorf("failed to create new listener: %w", err)
    89  	}
    90  
    91  	ctx, cancel := context.WithTimeout(ctx, defaultRuntimeDialTimeout)
    92  	defer cancel()
    93  	err = client.Initialize(ctx, version.RPCChainVMProtocol, listener.Addr().String())
    94  	if err != nil {
    95  		_ = listener.Close()
    96  		return fmt.Errorf("failed to initialize vm runtime: %w", err)
    97  	}
    98  
    99  	// start RPC Chain VM server
   100  	grpcutils.Serve(listener, server)
   101  
   102  	return nil
   103  }
   104  
   105  // Returns an RPC Chain VM server serving health and VM services.
   106  func newVMServer(vm block.ChainVM, allowShutdown *utils.Atomic[bool], opts ...grpcutils.ServerOption) *grpc.Server {
   107  	server := grpcutils.NewServer(opts...)
   108  	vmpb.RegisterVMServer(server, NewServer(vm, allowShutdown))
   109  
   110  	health := health.NewServer()
   111  	health.SetServingStatus("", healthpb.HealthCheckResponse_SERVING)
   112  	healthpb.RegisterHealthServer(server, health)
   113  
   114  	return server
   115  }