main.go

// Copyright (c) 2019 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"syscall"
	"time"

	"go.uber.org/yarpc"
	apipeer "go.uber.org/yarpc/api/peer"
	"go.uber.org/yarpc/internal/examples/thrift-hello/hello/echo"
	"go.uber.org/yarpc/internal/examples/thrift-hello/hello/echo/helloclient"
	"go.uber.org/yarpc/internal/examples/thrift-hello/hello/echo/helloserver"
	"go.uber.org/yarpc/peer"
	"go.uber.org/yarpc/peer/hostport"
	"go.uber.org/yarpc/peer/pendingheap"
	"go.uber.org/yarpc/peer/roundrobin"
	"go.uber.org/yarpc/transport/http"
)

var (
	flagSet   = flag.NewFlagSet("thrift-hello", flag.ExitOnError)
	flagWait  = flagSet.Bool("wait", false, "Wait for a signal to exit")
	flagRR    = flagSet.Bool("round-robin", false, "Use round-robin instead of fewest-pending-requests load balancer")
	flagCount = flagSet.Int("count", 1, "How many requests to send, or -1 to run indefinitely")
)

func main() {
	if err := do(); err != nil {
		log.Fatal(err)
	}
}

func do() error {
	if err := flagSet.Parse(os.Args[1:]); err != nil {
		return err
	}
	// Configure a YARPC dispatcher for the service "hello",
	// expose the service over an HTTP inbound on port 8086,
	// and configure outbound calls to service "hello" over HTTP port 8086 as well
	http := http.NewTransport()

	// Choose a load balancer (peer list)
	var pl apipeer.ChooserList
	if *flagRR {
		pl = roundrobin.New(http)
	} else {
		pl = pendingheap.New(http)
	}

	// Bind the peer list to a peer list updater (static peers), making the
	// peer chooser.
	pc := peer.Bind(pl, peer.BindPeers([]apipeer.Identifier{
		hostport.Identify("127.0.0.1:8086"),
		hostport.Identify("127.0.0.2:8086"),
		hostport.Identify("127.0.0.3:8086"),
	}))

	dispatcher := yarpc.NewDispatcher(yarpc.Config{
		Name: "hello",
		Inbounds: yarpc.Inbounds{
			http.NewInbound(":8086"),
		},
		Outbounds: yarpc.Outbounds{
			"hello": {
				Unary: http.NewOutbound(pc),
			},
		},
	})

	// register the Thrift handler which implements echo.thrift,
	// whose shapes are available in the hello/ dir, as generated by
	// the go:generate statement at the top of this file
	dispatcher.Register(helloserver.New(&helloHandler{}))

	// start the dispatcher, which enables requests to be sent and received
	if err := dispatcher.Start(); err != nil {
		return err
	}
	defer dispatcher.Stop()

	// create a Thrift client configured to call the "hello" service
	// using the dispatcher and associated "hello" outbound
	client := helloclient.New(dispatcher.ClientConfig("hello"))

	for i := 0; i != *flagCount; i++ {
		// build a context with a 1 second deadline
		ctx, cancel := context.WithTimeout(context.Background(), time.Second)

		// use the Thrift client to call the Echo procedure using YARPC
		res, err := client.Echo(ctx, &echo.EchoRequest{Message: "Hello world", Count: 1})
		if err != nil {
			fmt.Println(err)
		} else {
			fmt.Println(res)
		}

		cancel()
		time.Sleep(100 * time.Millisecond)
	}

	if *flagWait {
		// Gracefully shut down if we receive an interrupt (^C) or a kill signal.
		// Upon returning, this will unravel cancel() to abort the outbound request
		// and dispatcher.Stop() which will block until graceful shutdown.
		signals := make(chan os.Signal, 1)
		signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)
		<-signals
	}
	return nil
}

type helloHandler struct{}

func (h helloHandler) Echo(ctx context.Context, e *echo.EchoRequest) (*echo.EchoResponse, error) {
	return &echo.EchoResponse{Message: e.Message, Count: e.Count + 1}, nil
}