From 78101ac55ecc1b290fc406e1cbdf54bb76851f21 Mon Sep 17 00:00:00 2001
From: Wojciech Tyczynski <wojtekt@google.com>
Date: Mon, 18 Jul 2016 10:04:36 +0200
Subject: [PATCH] Fix crashes in schedulercache

---
 plugin/pkg/scheduler/schedulercache/cache.go  |  1 +
 .../scheduler/schedulercache/cache_test.go    | 30 +++++++++++++++++++
 2 files changed, 31 insertions(+)

diff --git a/plugin/pkg/scheduler/schedulercache/cache.go b/plugin/pkg/scheduler/schedulercache/cache.go
index acf2edef607..556aea30826 100644
--- a/plugin/pkg/scheduler/schedulercache/cache.go
+++ b/plugin/pkg/scheduler/schedulercache/cache.go
@@ -149,6 +149,7 @@ func (cache *schedulerCache) ForgetPod(pod *api.Pod) error {
 		if err != nil {
 			return err
 		}
+		delete(cache.assumedPods, key)
 		delete(cache.podStates, key)
 	default:
 		return fmt.Errorf("pod state wasn't assumed but get forgotten. Pod key: %v", key)
diff --git a/plugin/pkg/scheduler/schedulercache/cache_test.go b/plugin/pkg/scheduler/schedulercache/cache_test.go
index d8bc4a98467..fb0e062f9ce 100644
--- a/plugin/pkg/scheduler/schedulercache/cache_test.go
+++ b/plugin/pkg/scheduler/schedulercache/cache_test.go
@@ -457,6 +457,36 @@ func TestRemovePod(t *testing.T) {
 	}
 }
 
+func TestForgetPod(t *testing.T) {
+	nodeName := "node"
+	basePod := makeBasePod(nodeName, "test", "100m", "500", []api.ContainerPort{{HostPort: 80}})
+	tests := []struct {
+		pods []*api.Pod
+	}{{
+		pods: []*api.Pod{basePod},
+	}}
+	now := time.Now()
+	ttl := 10 * time.Second
+
+	for i, tt := range tests {
+		cache := newSchedulerCache(ttl, time.Second, nil)
+		for _, pod := range tt.pods {
+			if err := cache.assumePod(pod, now); err != nil {
+				t.Fatalf("assumePod failed: %v", err)
+			}
+		}
+		for _, pod := range tt.pods {
+			if err := cache.ForgetPod(pod); err != nil {
+				t.Fatalf("ForgetPod failed: %v", err)
+			}
+		}
+		cache.cleanupAssumedPods(now.Add(2 * ttl))
+		if n := cache.nodes[nodeName]; n != nil {
+			t.Errorf("#%d: expecting pod deleted and nil node info, get=%s", i, n)
+		}
+	}
+}
+
 func BenchmarkList1kNodes30kPods(b *testing.B) {
 	cache := setupCacheOf1kNodes30kPods(b)
 	b.ResetTimer()