Refactor metadataserver to enable optimistic IAM fetch

cmd/titus-metadata-service/main.go

@@ -2,17 +2,19 @@ package main
 
 import (
 	"context"
-	"flag"
 	"net"
 	"net/http"
 	"os"
 	"strconv"
 	"syscall"
+	"time"
 
 	"github.com/Netflix/titus-executor/logsutil"
 	"github.com/Netflix/titus-executor/metadataserver"
+	"github.com/Netflix/titus-executor/metadataserver/types"
 	log "github.com/sirupsen/logrus"
 	"golang.org/x/sys/unix"
+	"gopkg.in/urfave/cli.v1"
 )
 
 // 169 is the first octet of 169.254...
@@ -49,47 +51,91 @@ func makeFDListener(fd int64) net.Listener {
 	return l
 }
 
-func getEnv(key string) string {
-	val := os.Getenv(key)
-	if val == "" {
-		log.WithField("key", key).Fatal("Expected environmental variable unset: ", key)
-	}
-	return val
-
-}
-
 func main() {
+	app := cli.NewApp()
+	app.Name = "titus-metadata-service"
 	var listenerFd int64
 	var listenPort int
 	var debug bool
 	var backingMetadataServer string
-
-	flag.StringVar(&backingMetadataServer, "backing-metadata-server", "http://169.254.169.254/", "The URI of the AWS metadata server you want to use")
-	flag.Int64Var(&listenerFd, "listener-fd", -1, "Use a specific fd for listening on")
-	flag.IntVar(&listenPort, "listener-port", defaultListeningPort, "Use specific port to listen on")
-	flag.BoolVar(&debug, "debug", false, "Set to true to debug logging")
-
-	flag.Parse()
-	if debug {
-		log.SetLevel(log.DebugLevel)
-	} else {
-		log.SetLevel(log.InfoLevel)
+	var optimistic bool
+	var region string
+	var iamARN string
+	var titusTaskInstanceID string
+	var ipv4Address string
+	app.Flags = []cli.Flag{
+		cli.StringFlag{
+			Name:        "backing-metadata-server",
+			Value:       "http://169.254.169.254/",
+			Usage:       "The URI of the AWS metadata server you want to use",
+			Destination: &backingMetadataServer,
+		},
+		cli.Int64Flag{
+			Name:        "listener-fd",
+			Value:       -1,
+			Usage:       "Use a specific fd for listening on",
+			Destination: &listenerFd,
+		},
+		cli.BoolFlag{
+			Name:        "debug",
+			Usage:       "Set to true to enable debug logging",
+			Destination: &debug,
+		},
+		cli.IntFlag{
+			Name:        "listener-port",
+			Value:       defaultListeningPort,
+			Usage:       "Use specific port to listen on",
+			Destination: &listenPort,
+		},
+		cli.BoolFlag{
+			Name:        "optimistic",
+			Usage:       "If you set this to to true, the IAM service will optimistically fetch IAM credentials",
+			Destination: &optimistic,
+			EnvVar:      types.TitusOptimisticIAMVariableName,
+		},
+		cli.StringFlag{
+			Name:        "region",
+			Usage:       "The STS service region to use",
+			Destination: &region,
+			Value:       "",
+			EnvVar:      "EC2_REGION",
+		},
+		cli.StringFlag{
+			Name:        "iam-role",
+			EnvVar:      "TITUS_IAM_ROLE",
+			Destination: &iamARN,
+		},
+		cli.StringFlag{
+			Name:        "titus-task-instance-id",
+			EnvVar:      "TITUS_TASK_INSTANCE_ID",
+			Destination: &titusTaskInstanceID,
+		},
+		cli.StringFlag{
+			Name:        "ipv4-address",
+			EnvVar:      "EC2_LOCAL_IPV4",
+			Destination: &ipv4Address,
+		},
 	}
-	logsutil.MaybeSetupLoggerIfOnJournaldAvailable()
-
-	/* Get the requisite configuration from environment variables */
-	iamARN := getEnv("TITUS_IAM_ROLE")
-	titusTaskInstanceID := getEnv("TITUS_TASK_INSTANCE_ID")
-	ipv4Address := getEnv("EC2_LOCAL_IPV4")
-
-	listener := getListener(listenPort, listenerFd)
-	ms, err := metadataserver.NewMetaDataServer(context.Background(), backingMetadataServer, iamARN, titusTaskInstanceID, ipv4Address)
-	if err != nil {
-		log.WithError(err).Fatal("Could not start / create metadataserver")
+	app.Action = func(c *cli.Context) error {
+		if debug {
+			log.SetLevel(log.DebugLevel)
+		} else {
+			log.SetLevel(log.InfoLevel)
+		}
+		logsutil.MaybeSetupLoggerIfOnJournaldAvailable()
+
+		/* Get the requisite configuration from environment variables */
+		listener := getListener(listenPort, listenerFd)
+		ms := metadataserver.NewMetaDataServer(context.Background(), backingMetadataServer, iamARN, titusTaskInstanceID, ipv4Address, region, optimistic)
+		go notifySystemd()
+		if err := http.Serve(listener, ms); err != nil {
+			return err
+		}
+		log.Info("Done serving?")
+		time.Sleep(1 * time.Second)
+		return nil
 	}
-	go notifySystemd()
-	if err := http.Serve(listener, ms); err != nil {
-		log.Fatal(err)
+	if err := app.Run(os.Args); err != nil {
+		log.WithError(err).Fatal()
 	}
-
 }

executor/runtime/docker/docker.go

@@ -34,6 +34,7 @@ import (
 	docker "github.com/docker/docker/client"
 	"github.com/docker/go-units"
 	"github.com/ftrvxmtrx/fd"
+	metadataserverTypes "github.com/Netflix/titus-executor/metadataserver/types"
 	"github.com/hashicorp/go-multierror"
 	"github.com/pkg/errors"
 	log "github.com/sirupsen/logrus"
@@ -975,6 +976,17 @@ func (r *DockerRuntime) createTitusEnvironmentFile(c *runtimeTypes.Container) er
 		return os.Remove(envFile)
 	})
 
+	if _, err := f.WriteString(fmt.Sprintf("TITUS_EXECUTOR_DIR=%s\n", strconv.QuoteToASCII(os.Args[0]))); err != nil {
+		return err
+	}
+
+	if optimisticTokenFetch, parseErr := c.GetOptimisticIAMTokenFetch(); parseErr != nil {
+		return parseErr
+	} else if optimisticTokenFetch {
+		if _, err := f.WriteString(fmt.Sprintf("%s=true\n", metadataserverTypes.TitusOptimisticIAMVariableName)); err != nil {
+			return err
+		}
+	}
 	/* writeTitusEnvironmentFile closes the file for us */
 	return writeTitusEnvironmentFile(c.Env, f)
 }
@@ -994,10 +1006,6 @@ func writeTitusEnvironmentFile(env map[string]string, w io.Writer) error {
 		}
 	}
 
-	if _, err := writer.WriteString(fmt.Sprintf("TITUS_EXECUTOR_DIR=%s\n", strconv.QuoteToASCII(os.Args[0]))); err != nil {
-		return err
-	}
-
 	return writer.Flush()
 }
 

executor/runtime/types/types.go

@@ -26,9 +26,10 @@ import (
 const (
 	hostnameStyleParam = "titusParameter.agent.hostnameStyle"
 	// FuseEnabledParam is a container atttribute set to enable FUSE
-	FuseEnabledParam       = "titusParameter.agent.fuseEnabled"
-	assignIPv6AddressParam = "titusParameter.agent.assignIPv6Address"
-	ttyEnabledParam        = "titusParameter.agent.ttyEnabled"
+	FuseEnabledParam             = "titusParameter.agent.fuseEnabled"
+	assignIPv6AddressParam       = "titusParameter.agent.assignIPv6Address"
+	ttyEnabledParam              = "titusParameter.agent.ttyEnabled"
+	optimisticIAMTokenFetchParam = "titusParameter.agent.optimisticIAMTokenFetch"
 )
 
 const (
@@ -380,6 +381,15 @@ func (c *Container) GetKeepLocalFileAfterUpload() (bool, error) {
 	return strconv.ParseBool(keepLocalFileAfterUploadStr)
 }
 
+// GetOptimisticIAMTokenFetch indicates whether or not we should delete log files after uploading them
+func (c *Container) GetOptimisticIAMTokenFetch() (bool, error) {
+	optimisticIAMTokenFetchStr, ok := c.TitusInfo.GetPassthroughAttributes()[optimisticIAMTokenFetchParam]
+	if !ok {
+		return false, nil
+	}
+	return strconv.ParseBool(optimisticIAMTokenFetchStr)
+}
+
 // Resources specify constraints to be applied to a Container
 type Resources struct {
 	Mem       int64 // in MiB