Adding support for simple-tcp service (#4)

README.md

@@ -26,7 +26,7 @@ This flag is passed as a map like `--hosts-ping example.com=pingServer.example.c
 
 This command can be used to run a bunch of simultaneous allocations and connections. See the help for configuration.
 
-NOTE: This currently only supports the Agones simple-udp server. It makes a connect, says hello, waits, and then says goodbye and EXIT.
+NOTE: This currently only supports the Agones simple-udp or simple-tcp server. It makes a connect, says hello, waits, and then says goodbye and EXIT.
 
 ## Attribution
 

cmd/root.go

@@ -46,6 +46,7 @@ var (
 	labelSelector map[string]string
 	pingTargets   []string
 	maxRetries    int
+	protocol      string
 )
 
 func init() {
@@ -54,7 +55,7 @@ func init() {
 	rootCmd.PersistentFlags().StringVar(&caCertFile, "ca-cert", "", "The path the CA cert file in PEM format")
 	rootCmd.PersistentFlags().StringSliceVar(&hosts, "hosts", nil, "A list of possible allocation servers. If nil, you must set hosts-ping")
 	rootCmd.PersistentFlags().StringToStringVar(&pingServers, "hosts-ping", nil, "A map hosts and and ping servers. If nil, you must set hosts.")
-	rootCmd.PersistentFlags().StringVarP(&namespace, "namespace", "n", "", "The namespace of gameservers to request from")
+	rootCmd.PersistentFlags().StringVarP(&namespace, "namespace", "n", "default", "The namespace of gameservers to request from")
 	rootCmd.PersistentFlags().BoolVarP(&multicluster, "multicluster", "m", false, "If true, multicluster allocation will be requested")
 	rootCmd.PersistentFlags().StringToStringVar(&labelSelector, "labels", nil, "A map of labels to match on the allocation.")
 	rootCmd.PersistentFlags().IntVar(&maxRetries, "max-retries", 10, "The maximum number of times to retry allocations.")
@@ -65,6 +66,7 @@ func init() {
 	loadTestCmd.PersistentFlags().IntVarP(&demoCount, "count", "c", 10, "The number of connections to make during the demo.")
 	loadTestCmd.PersistentFlags().IntVar(&demoDelay, "delay", 2, "The number of seconds to wait between connections")
 	loadTestCmd.PersistentFlags().IntVarP(&demoDuration, "duration", "d", 10, "The number of seconds to leave each connection open.")
+	loadTestCmd.PersistentFlags().StringVar(&protocol, "protocol", "udp", "The gameserver protocol. Either tcp or udp")
 
 	rootCmd.AddCommand(pingTestCmd)
 	pingTestCmd.PersistentFlags().StringSliceVarP(&pingTargets, "targets", "t", nil, "The list of targets to ping.")
@@ -141,7 +143,7 @@ var loadTestCmd = &cobra.Command{
 		if err != nil {
 			klog.Fatal(err)
 		}
-		err = allocatorClient.RunUDPLoad(demoCount, demoDelay, demoDuration)
+		err = allocatorClient.RunLoad(demoCount, demoDelay, demoDuration, protocol)
 		if err != nil {
 			klog.Fatal(err)
 		}
@@ -156,6 +158,9 @@ var pingTestCmd = &cobra.Command{
 		if pingTargets == nil {
 			return fmt.Errorf("You must pass a list of target hostanmes or IP addresses")
 		}
+		if protocol != "udp" && protocol != "tcp" {
+			return fmt.Errorf("You must specify a gameserver protocol using --protocol that is either udp or tcp")
+		}
 		return nil
 	},
 	Run: func(cmd *cobra.Command, args []string) {

pkg/allocator/connect.go

@@ -0,0 +1,121 @@
+/*
+Copyright 2020 Fairwinds
+
+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 allocator
+
+import (
+	"bufio"
+	"fmt"
+	"net"
+	"sync"
+	"time"
+
+	"k8s.io/klog"
+)
+
+// RunLoad runs many concurrent game connections on a simple UDP or TCP server
+// This is designed to test the allocator service and autoscaling of the game servers.
+func (c *Client) RunLoad(count int, delay int, duration int, proto string) error {
+	var wg sync.WaitGroup
+	for i := 0; i < count; i++ {
+		wg.Add(1)
+		go c.testConnection(i, &wg, duration, proto)
+		time.Sleep(time.Duration(delay) * time.Second)
+	}
+	wg.Wait()
+	return nil
+}
+
+func (c *Client) testConnection(id int, wg *sync.WaitGroup, duration int, proto string) {
+	defer wg.Done()
+
+	a, err := c.AllocateGameserverWithRetry()
+	if err != nil {
+		klog.Error(err.Error())
+		return
+	}
+
+	klog.V(3).Infof("%d - got allocation %s %d. Proceeding to connection...\n", id, a.Address, a.Port)
+	err = a.testConnection(id, duration, proto)
+	if err != nil {
+		klog.Error(err)
+	}
+}
+
+// testConnection tests a series of connections to the simple-udp server gameserver example
+func (a *Allocation) testConnection(id int, duration int, proto string) error {
+	endpoint := fmt.Sprintf("%s:%d", a.Address, a.Port)
+
+	switch proto {
+	case "tcp":
+		conn, err := net.Dial("tcp", endpoint)
+		if err != nil {
+			return err
+		}
+		klog.V(2).Infof("%d - connected to gameserver and sending hello", id)
+		fmt.Fprintf(conn, "HELLO\n")
+		status, _ := bufio.NewReader(conn).ReadString('\n')
+		klog.V(3).Infof("%d - response: %s", id, status)
+
+		klog.V(3).Infof("%d - sleeping %d seconds to view logs", id, duration)
+		time.Sleep(time.Duration(duration) * time.Second)
+
+		klog.V(3).Infof("%d - closing connection", id)
+		fmt.Fprintf(conn, "EXIT\n")
+		return nil
+
+	case "udp":
+		conn, err := net.ListenPacket(proto, ":0")
+		if err != nil {
+			return err
+		}
+		defer conn.Close()
+
+		dst, err := net.ResolveUDPAddr(proto, endpoint)
+		if err != nil {
+			return err
+		}
+
+		klog.V(2).Infof("%d - connected to gameserver and sending hello", id)
+
+		// Hello
+		msg := fmt.Sprintf("Hello from process %d!", id)
+		_, err = conn.WriteTo([]byte(msg), dst)
+		if err != nil {
+			return err
+		}
+
+		// Wait
+		klog.V(3).Infof("%d - sleeping %d seconds to view logs", id, duration)
+		time.Sleep(time.Duration(duration) * time.Second)
+
+		// Goodbye
+		msg = fmt.Sprintf("Goodbye from process %d.", id)
+		_, err = conn.WriteTo([]byte(msg), dst)
+		if err != nil {
+			return err
+		}
+
+		klog.V(3).Infof("%d - closing connection", id)
+		_, err = conn.WriteTo([]byte("EXIT"), dst)
+		if err != nil {
+			return err
+		}
+		return nil
+	default:
+		return fmt.Errorf("proto must be one of (udp|tcp)")
+	}
+}