agones.dev/agones@v1.53.0/test/load/allocation/allocationload/allocationload.go (about)

     1  // Copyright 2020 Google LLC All Rights Reserved.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License
    14  
    15  //nolint:typecheck
    16  package main
    17  
    18  import (
    19  	"context"
    20  	"crypto/tls"
    21  	"crypto/x509"
    22  	"flag"
    23  	"fmt"
    24  	"os"
    25  	"sync"
    26  	"time"
    27  
    28  	pb "agones.dev/agones/pkg/allocation/go"
    29  	"github.com/pkg/errors"
    30  	"google.golang.org/grpc"
    31  	"google.golang.org/grpc/credentials"
    32  )
    33  
    34  func main() {
    35  	keyFile := flag.String("key", "missing key", "the private key file for the client certificate in PEM format")
    36  	certFile := flag.String("cert", "missing cert", "the public key file for the client certificate in PEM format")
    37  	cacertFile := flag.String("cacert", "missing cacert", "the CA cert file for server signing certificate in PEM format")
    38  	externalIP := flag.String("ip", "missing external IP", "the external IP for allocator server")
    39  	port := flag.String("port", "443", "the port for allocator server")
    40  	namespace := flag.String("namespace", "default", "the game server kubernetes namespace")
    41  	multicluster := flag.Bool("multicluster", false, "set to true to enable the multi-cluster allocation")
    42  	numOfClients := flag.Int("numberofclients", 1, "number of clients to do allocations in parallel")
    43  	perClientAllocs := flag.Int("perclientallocations", 1, "number of allocations to be done per client")
    44  
    45  	flag.Parse()
    46  
    47  	endpoint := *externalIP + ":" + *port
    48  	cert, err := os.ReadFile(*certFile)
    49  	if err != nil {
    50  		panic(err)
    51  	}
    52  	key, err := os.ReadFile(*keyFile)
    53  	if err != nil {
    54  		panic(err)
    55  	}
    56  	cacert, err := os.ReadFile(*cacertFile)
    57  	if err != nil {
    58  		panic(err)
    59  	}
    60  
    61  	request := &pb.AllocationRequest{
    62  		Namespace: *namespace,
    63  		MultiClusterSetting: &pb.MultiClusterSetting{
    64  			Enabled: *multicluster,
    65  		},
    66  	}
    67  
    68  	dialOpts, err := createRemoteClusterDialOption(cert, key, cacert)
    69  	if err != nil {
    70  		panic(err)
    71  	}
    72  
    73  	fmt.Printf("started: %v\n", time.Now())
    74  
    75  	var wg sync.WaitGroup
    76  
    77  	// Allocate GS by numOfClients in parallel
    78  	for k := 0; k < *numOfClients; k++ {
    79  		wg.Add(1)
    80  
    81  		go func(clientID int) {
    82  			defer wg.Done()
    83  			conn, err := grpc.NewClient(endpoint, dialOpts)
    84  			if err != nil {
    85  				fmt.Printf("(failed(client=%v) to get connection: %v\n", clientID, err)
    86  				return
    87  			}
    88  			grpcClient := pb.NewAllocationServiceClient(conn)
    89  
    90  			// Create a WaitGroup for the allocations within this client
    91  			var allocWg sync.WaitGroup
    92  
    93  			for i := 0; i < *perClientAllocs; i++ {
    94  				allocWg.Add(1)
    95  
    96  				// Launch each allocation in its own goroutine
    97  				go func(allocID int) {
    98  					defer allocWg.Done()
    99  					_, err := grpcClient.Allocate(context.Background(), request)
   100  					if err != nil {
   101  						fmt.Printf("(failed(client=%v,allocation=%v): %v\n", clientID, allocID+1, err)
   102  					}
   103  				}(i)
   104  			}
   105  
   106  			// Wait for all allocations from this client to complete
   107  			allocWg.Wait()
   108  			_ = conn.Close()
   109  		}(k)
   110  	}
   111  	wg.Wait()
   112  	fmt.Printf("finished: %v\n", time.Now())
   113  }
   114  
   115  // createRemoteClusterDialOption creates a grpc client dial option with TLS configuration.
   116  func createRemoteClusterDialOption(clientCert, clientKey, caCert []byte) (grpc.DialOption, error) {
   117  	// Load client cert
   118  	cert, err := tls.X509KeyPair(clientCert, clientKey)
   119  	if err != nil {
   120  		return nil, err
   121  	}
   122  
   123  	tlsConfig := &tls.Config{Certificates: []tls.Certificate{cert}}
   124  	if len(caCert) != 0 {
   125  		// Load CA cert, if provided and trust the server certificate.
   126  		// This is required for self-signed certs.
   127  		tlsConfig.RootCAs = x509.NewCertPool()
   128  		if !tlsConfig.RootCAs.AppendCertsFromPEM(caCert) {
   129  			return nil, errors.New("only PEM format is accepted for server CA")
   130  		}
   131  	}
   132  
   133  	return grpc.WithTransportCredentials(credentials.NewTLS(tlsConfig)), nil
   134  }