Kill the store

cmd/k8s.go

@@ -1,43 +1,70 @@
 package cmd
 
 import (
+	"errors"
+	"fmt"
 	"time"
 
 	"k8s.io/kubernetes/pkg/api"
 	kcache "k8s.io/kubernetes/pkg/client/cache"
 	"k8s.io/kubernetes/pkg/client/restclient"
 	client "k8s.io/kubernetes/pkg/client/unversioned"
+	"k8s.io/kubernetes/pkg/controller"
 	selector "k8s.io/kubernetes/pkg/fields"
 	"k8s.io/kubernetes/pkg/util/wait"
 )
 
 const (
 	// Resync period for the kube controller loop.
-	resyncPeriod = 30 * time.Minute
+	resyncPeriod   = 30 * time.Minute
+	podIPIndexName = "byPodIP"
 )
 
+var errPodNotFound = errors.New("Pod with specified IP not found")
+
 type k8s struct {
 	*client.Client
+	podStore      kcache.Indexer
+	podController *kcache.Controller
+}
+
+func podIPIndexFunc(obj interface{}) ([]string, error) {
+	pod, ok := obj.(*api.Pod)
+	if !ok {
+		return nil, fmt.Errorf("obj not pod: %+v", obj)
+	}
+	if pod.Status.PodIP != "" && controller.IsPodActive(pod) {
+		return []string{pod.Status.PodIP}, nil
+	}
+	return nil, nil
 }
 
 // Returns a cache.ListWatch that gets all changes to pods.
 func (k8s *k8s) createPodLW() *kcache.ListWatch {
 	return kcache.NewListWatchFromClient(k8s, "pods", api.NamespaceAll, selector.Everything())
 }
 
-func (k8s *k8s) watchForPods(podManager kcache.ResourceEventHandler) kcache.Store {
-	podStore, podController := kcache.NewInformer(
-		k8s.createPodLW(),
-		&api.Pod{},
-		resyncPeriod,
-		kcache.ResourceEventHandlerFuncs{
-			AddFunc:    podManager.OnAdd,
-			DeleteFunc: podManager.OnDelete,
-			UpdateFunc: podManager.OnUpdate,
-		},
-	)
-	go podController.Run(wait.NeverStop)
-	return podStore
+func (k8s *k8s) PodByIP(IP string) (*api.Pod, error) {
+	pods, err := k8s.podStore.ByIndex(podIPIndexName, IP)
+	if err != nil {
+		return nil, err
+	}
+
+	if len(pods) == 0 {
+		return nil, errPodNotFound
+	}
+	if len(pods) == 1 {
+		return pods[0].(*api.Pod), nil
+	}
+	podNames := make([]string, len(pods))
+	for i, pod := range pods {
+		podNames[i] = pod.(*api.Pod).ObjectMeta.Name
+	}
+	return nil, fmt.Errorf("%d pods (%v) with the ip %s indexed", len(pods), podNames, IP)
+}
+
+func (k8s *k8s) Run() {
+	k8s.podController.Run(wait.NeverStop)
 }
 
 func newK8s(host, token string, insecure bool) (*k8s, error) {
@@ -56,5 +83,13 @@ func newK8s(host, token string, insecure bool) (*k8s, error) {
 	if err != nil {
 		return nil, err
 	}
-	return &k8s{c}, nil
+	k8s := &k8s{c, nil, nil}
+	k8s.podStore, k8s.podController = kcache.NewIndexerInformer(
+		k8s.createPodLW(),
+		&api.Pod{},
+		resyncPeriod,
+		kcache.ResourceEventHandlerFuncs{},
+		kcache.Indexers{podIPIndexName: podIPIndexFunc},
+	)
+	return k8s, nil
 }

cmd/server.go

@@ -6,6 +6,7 @@ import (
 	"net/http"
 	"net/http/httputil"
 	"strings"
+	"time"
 
 	log "github.com/Sirupsen/logrus"
 	"github.com/cenk/backoff"
@@ -30,7 +31,6 @@ type Server struct {
 	Version         bool
 	iam             *iam
 	k8s             *k8s
-	store           *store
 }
 
 type appHandler func(http.ResponseWriter, *http.Request)
@@ -57,18 +57,28 @@ func parseRemoteAddr(addr string) string {
 
 func (s *Server) getRole(IP string) (string, error) {
 	var role string
-	var err error
 	operation := func() error {
-		role, err = s.store.Get(IP)
-		return err
+		pod, err := s.k8s.PodByIP(IP)
+		if err != nil {
+			if err != errPodNotFound {
+				log.Warnf("Role request from IP %s resulted in error: %s", IP, err.Error())
+			}
+			return err
+		}
+		role = pod.Annotations[s.IAMRoleKey]
+		return nil
 	}
 
-	err = backoff.Retry(operation, backoff.NewExponentialBackOff())
-	if err != nil {
-		return "", err
+	exponential := backoff.NewExponentialBackOff()
+	exponential.MaxElapsedTime = 1 * time.Second
+	exponential.MaxInterval = 30 * time.Millisecond
+
+	err := backoff.Retry(operation, exponential)
+	if role == "" && err == nil {
+		err = errPodNotFound
 	}
 
-	return role, nil
+	return role, err
 }
 
 func (s *Server) securityCredentialsHandler(w http.ResponseWriter, r *http.Request) {
@@ -128,8 +138,7 @@ func (s *Server) Run(host, token string, insecure bool) error {
 		return err
 	}
 	s.k8s = k8s
-	s.store = newStore(s.IAMRoleKey, s.DefaultIAMRole)
-	s.k8s.watchForPods(s.store)
+	go k8s.Run()
 	s.iam = newIAM(s.BaseRoleARN)
 	r := mux.NewRouter()
 	r.Handle("/{version}/meta-data/iam/security-credentials/", appHandler(s.securityCredentialsHandler))