From f26325c988ddb9afd54153433c3520796df6269e Mon Sep 17 00:00:00 2001
From: Harshavardhana <harsha@minio.io>
Date: Wed, 15 Aug 2018 16:35:21 -0700
Subject: [PATCH] Support supplying custom drives per set count (#6261)

---
 cmd/endpoint-ellipses.go      | 41 +++++++++++++++++++++---
 cmd/endpoint-ellipses_test.go | 60 +++++++++++++++++++++++++++++++++++
 cmd/ui-errors.go              |  6 ++++
 3 files changed, 103 insertions(+), 4 deletions(-)

diff --git a/cmd/endpoint-ellipses.go b/cmd/endpoint-ellipses.go
index 3bdb4b386..0f3d67fc4 100644
--- a/cmd/endpoint-ellipses.go
+++ b/cmd/endpoint-ellipses.go
@@ -18,6 +18,8 @@ package cmd
 
 import (
 	"fmt"
+	"os"
+	"strconv"
 	"strings"
 
 	"github.com/minio/minio-go/pkg/set"
@@ -69,10 +71,26 @@ func getSetIndexes(args []string, totalSizes []uint64) (setIndexes [][]uint64, e
 		return nil, errInvalidArgument
 	}
 
+	// isValidSetSize - checks whether given count is a valid set size for erasure coding.
+	isValidSetSize := func(count uint64) bool {
+		return (count >= setSizes[0] && count <= setSizes[len(setSizes)-1] && count%2 == 0)
+	}
+
+	var customSetDriveCount uint64
+	if v := os.Getenv("MINIO_ERASURE_SET_DRIVE_COUNT"); v != "" {
+		customSetDriveCount, err = strconv.ParseUint(v, 10, 64)
+		if err != nil {
+			return nil, uiErrInvalidErasureSetSize(err)
+		}
+		if !isValidSetSize(customSetDriveCount) {
+			return nil, uiErrInvalidErasureSetSize(nil)
+		}
+	}
+
 	setIndexes = make([][]uint64, len(totalSizes))
 	for _, totalSize := range totalSizes {
 		// Check if totalSize has minimum range upto setSize
-		if totalSize < setSizes[0] {
+		if totalSize < setSizes[0] || totalSize < customSetDriveCount {
 			return nil, uiErrInvalidNumberOfErasureEndpoints(nil)
 		}
 	}
@@ -95,9 +113,24 @@ func getSetIndexes(args []string, totalSizes []uint64) (setIndexes [][]uint64, e
 		setSize = commonSize
 	}
 
-	// isValidSetSize - checks whether given count is a valid set size for erasure coding.
-	isValidSetSize := func(count uint64) bool {
-		return (count >= setSizes[0] && count <= setSizes[len(setSizes)-1] && count%2 == 0)
+	possibleSetCounts := func(setSize uint64) (ss []uint64) {
+		for _, s := range setSizes {
+			if setSize%s == 0 {
+				ss = append(ss, s)
+			}
+		}
+		return ss
+	}
+
+	if customSetDriveCount > 0 {
+		msg := fmt.Sprintf("Invalid set drive count, leads to non-uniform distribution for the given number of disks. Possible values for custom set count are %d", possibleSetCounts(setSize))
+		if customSetDriveCount > setSize {
+			return nil, uiErrInvalidErasureSetSize(nil).Msg(msg)
+		}
+		if setSize%customSetDriveCount != 0 {
+			return nil, uiErrInvalidErasureSetSize(nil).Msg(msg)
+		}
+		setSize = customSetDriveCount
 	}
 
 	// Check whether setSize is with the supported range.
diff --git a/cmd/endpoint-ellipses_test.go b/cmd/endpoint-ellipses_test.go
index dcbd0ca39..4379708ce 100644
--- a/cmd/endpoint-ellipses_test.go
+++ b/cmd/endpoint-ellipses_test.go
@@ -18,6 +18,7 @@ package cmd
 
 import (
 	"fmt"
+	"os"
 	"reflect"
 	"testing"
 
@@ -84,6 +85,65 @@ func TestGetDivisibleSize(t *testing.T) {
 	}
 }
 
+// Test tests calculating set indexes with ENV override for drive count.
+func TestGetSetIndexesEnvOverride(t *testing.T) {
+	testCases := []struct {
+		args        []string
+		totalSizes  []uint64
+		indexes     [][]uint64
+		envOverride string
+		success     bool
+	}{
+		{
+			[]string{"data{1...64}"},
+			[]uint64{64},
+			[][]uint64{{8, 8, 8, 8, 8, 8, 8, 8}},
+			"8",
+			true,
+		},
+		{
+			[]string{"data{1...60}"},
+			nil,
+			nil,
+			"8",
+			false,
+		},
+		{
+			[]string{"data{1...64}"},
+			nil,
+			nil,
+			"-1",
+			false,
+		},
+		{
+			[]string{"data{1...64}"},
+			nil,
+			nil,
+			"2",
+			false,
+		},
+	}
+
+	for i, testCase := range testCases {
+		t.Run(fmt.Sprintf("Test%d", i+1), func(t *testing.T) {
+			if err := os.Setenv("MINIO_ERASURE_SET_DRIVE_COUNT", testCase.envOverride); err != nil {
+				t.Fatal(err)
+			}
+			gotIndexes, err := getSetIndexes(testCase.args, testCase.totalSizes)
+			if err != nil && testCase.success {
+				t.Errorf("Expected success but failed instead %s", err)
+			}
+			if err == nil && !testCase.success {
+				t.Errorf("Expected failure but passed instead")
+			}
+			if !reflect.DeepEqual(testCase.indexes, gotIndexes) {
+				t.Errorf("Expected %v, got %v", testCase.indexes, gotIndexes)
+			}
+			os.Unsetenv("MINIO_ERASURE_SET_DRIVE_COUNT")
+		})
+	}
+}
+
 // Test tests calculating set indexes.
 func TestGetSetIndexes(t *testing.T) {
 	testCases := []struct {
diff --git a/cmd/ui-errors.go b/cmd/ui-errors.go
index 4a1928d79..b6f7d1ded 100644
--- a/cmd/ui-errors.go
+++ b/cmd/ui-errors.go
@@ -29,6 +29,12 @@ var (
 		"Browser can only accept `on` and `off` values. To disable web browser access, set this value to `off`",
 	)
 
+	uiErrInvalidErasureSetSize = newUIErrFn(
+		"Invalid erasure set size",
+		"Please check the passed value",
+		"Erasure set can only accept any of [4, 6, 8, 10, 12, 14, 16] values.",
+	)
+
 	uiErrInvalidWormValue = newUIErrFn(
 		"Invalid WORM value",
 		"Please check the passed value",