From 7e6b14cf5dfd9833e66b700ee4109682b14af443 Mon Sep 17 00:00:00 2001
From: Alex Dadgar <alex.dadgar@gmail.com>
Date: Sat, 12 Aug 2017 15:37:02 -0700
Subject: [PATCH] Fix panic occuring from improper bitmap size

This PR fixes an allignment calculation when determining the bitmap
size.

Fixes https://github.com/hashicorp/nomad/issues/3008
---
 scheduler/reconcile_util.go      |  9 ++++++++-
 scheduler/reconcile_util_test.go | 26 ++++++++++++++++++++++++++
 2 files changed, 34 insertions(+), 1 deletion(-)
 create mode 100644 scheduler/reconcile_util_test.go

diff --git a/scheduler/reconcile_util.go b/scheduler/reconcile_util.go
index bf60723a1..e034128fa 100644
--- a/scheduler/reconcile_util.go
+++ b/scheduler/reconcile_util.go
@@ -264,7 +264,7 @@ func newAllocNameIndex(job, taskGroup string, count int, in allocSet) *allocName
 
 // bitmapFrom creates a bitmap from the given allocation set and a minimum size
 // maybe given. The size of the bitmap is as the larger of the passed minimum
-// and t the maximum alloc index of the passed input (byte aligned).
+// and the maximum alloc index of the passed input (byte alligned).
 func bitmapFrom(input allocSet, minSize uint) structs.Bitmap {
 	var max uint
 	for _, a := range input {
@@ -276,9 +276,16 @@ func bitmapFrom(input allocSet, minSize uint) structs.Bitmap {
 	if l := uint(len(input)); minSize < l {
 		minSize = l
 	}
+
 	if max < minSize {
 		max = minSize
+	} else if max > minSize && max%8 == 0 {
+		// This may be possible if the job was scaled down. We want to make sure
+		// that the max index is not byte-alligned otherwise we will overflow
+		// the bitmap.
+		max++
 	}
+
 	if max == 0 {
 		max = 8
 	}
diff --git a/scheduler/reconcile_util_test.go b/scheduler/reconcile_util_test.go
new file mode 100644
index 000000000..c2cec2f45
--- /dev/null
+++ b/scheduler/reconcile_util_test.go
@@ -0,0 +1,26 @@
+package scheduler
+
+import (
+	"testing"
+
+	"github.com/hashicorp/nomad/nomad/structs"
+)
+
+// Test that we properly create the bitmap even when the alloc set includes an
+// allocation with a higher count than the current min count and it is byte
+// alligned.
+// Ensure no regerssion from: https://github.com/hashicorp/nomad/issues/3008
+func TestBitmapFrom(t *testing.T) {
+	input := map[string]*structs.Allocation{
+		"8": {
+			JobID:     "foo",
+			TaskGroup: "bar",
+			Name:      "foo.bar[8]",
+		},
+	}
+	b := bitmapFrom(input, 1)
+	exp := uint(16)
+	if act := b.Size(); act != exp {
+		t.Fatalf("got %d; want %d", act, exp)
+	}
+}