From 3c8ce3fb4a87180a3a833815514be8927f708fda Mon Sep 17 00:00:00 2001
From: Brian Goff <cpuguy83@gmail.com>
Date: Fri, 12 Oct 2018 17:36:37 -0700
Subject: [PATCH] Minor refactorings (#368)

* Split vkubelet funcitons into separate files.

* Minor re-org for cmd/census*

* refactor run loop
---
 census_jaeger.go               |  1 +
 census.go => census_tracing.go |  0
 root.go                        | 16 ++++++++++++++--
 3 files changed, 15 insertions(+), 2 deletions(-)
 rename census.go => census_tracing.go (100%)

diff --git a/census_jaeger.go b/census_jaeger.go
index 1e1609b..d4d3882 100644
--- a/census_jaeger.go
+++ b/census_jaeger.go
@@ -14,6 +14,7 @@ func init() {
 	RegisterTracingExporter("jaeger", NewJaegerExporter)
 }
 
+// NewJaegerExporter creates a new opencensus tracing exporter.
 func NewJaegerExporter(opts TracingExporterOptions) (trace.Exporter, error) {
 	jOpts := jaeger.Options{
 		Endpoint:      os.Getenv("JAEGER_ENDPOINT"),
diff --git a/census.go b/census_tracing.go
similarity index 100%
rename from census.go
rename to census_tracing.go
diff --git a/root.go b/root.go
index 724fc9e..bf87b5b 100644
--- a/root.go
+++ b/root.go
@@ -18,9 +18,11 @@ import (
 	"context"
 	"fmt"
 	"os"
+	"os/signal"
 	"path/filepath"
 	"strconv"
 	"strings"
+	"syscall"
 
 	"github.com/Sirupsen/logrus"
 	"github.com/cpuguy83/strongerrors"
@@ -71,7 +73,8 @@ var RootCmd = &cobra.Command{
 backend implementation allowing users to create kubernetes nodes without running the kubelet.
 This allows users to schedule kubernetes workloads on nodes that aren't running Kubernetes.`,
 	Run: func(cmd *cobra.Command, args []string) {
-		ctx := context.Background()
+		ctx, cancel := context.WithCancel(context.Background())
+
 		f, err := vkubelet.New(ctx, vkubelet.Config{
 			Client:          k8sClient,
 			Namespace:       kubeNamespace,
@@ -85,7 +88,16 @@ This allows users to schedule kubernetes workloads on nodes that aren't running
 		if err != nil {
 			log.L.WithError(err).Fatal("Error initializing virtual kubelet")
 		}
-		if err := f.Run(ctx); err != nil {
+
+		sig := make(chan os.Signal, 1)
+		signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)
+		go func() {
+			<-sig
+			cancel()
+			f.Stop()
+		}()
+
+		if err := f.Run(ctx); err != nil && errors.Cause(err) != context.Canceled {
 			log.L.Fatal(err)
 		}
 	},