From 398deb7cc1cf06ef9fd288438f48728bb2cd5c15 Mon Sep 17 00:00:00 2001
From: Balaji Subramaniam <balaji.subramaniam@intel.com>
Date: Tue, 15 Nov 2016 14:29:27 -0800
Subject: [PATCH] Enabled graceful failure if discovery for a source fails. -
 Handles errors from discovery of source. - Handles panics from discovery of
 source using recover(). - Added tests.

---
 main.go         | 15 ++++++++++++---
 main_test.go    | 15 +++++++++++++++
 sources.go      | 16 ----------------
 test_sources.go | 28 ++++++++++++++++++++++++++++
 4 files changed, 55 insertions(+), 19 deletions(-)
 create mode 100644 test_sources.go

diff --git a/main.go b/main.go
index ccaec5337..cac76f3fb 100644
--- a/main.go
+++ b/main.go
@@ -137,7 +137,9 @@ func main() {
 	for _, source := range sources {
 		labelsFromSource, err := getFeatureLabels(source)
 		if err != nil {
-			stderrLogger.Fatalf("discovery failed for source [%s]: %s", source.Name(), err.Error())
+			stderrLogger.Printf("discovery failed for source [%s]: %s", source.Name(), err.Error())
+			stderrLogger.Printf("continuing ...")
+			continue
 		}
 
 		for name, value := range labelsFromSource {
@@ -164,8 +166,15 @@ func main() {
 
 // getFeatureLabels returns node labels for features discovered by the
 // supplied source.
-func getFeatureLabels(source FeatureSource) (Labels, error) {
-	labels := Labels{}
+func getFeatureLabels(source FeatureSource) (labels Labels, err error) {
+	defer func() {
+		if r := recover(); r != nil {
+			stderrLogger.Printf("panic occured during discovery of source [%s]: %v", source.Name(), r)
+			err = fmt.Errorf("%v", r)
+		}
+	}()
+
+	labels = Labels{}
 	features, err := source.Discover()
 	if err != nil {
 		return nil, err
diff --git a/main_test.go b/main_test.go
index e92ced292..9571040db 100644
--- a/main_test.go
+++ b/main_test.go
@@ -165,3 +165,18 @@ func TestRemoveLabels(t *testing.T) {
 		})
 	})
 }
+
+func TestGetFeatureLabels(t *testing.T) {
+	Convey("When I get feature labels and panic occurs during discovery of a feature source", t, func() {
+		fakePanicFeatureSource := FeatureSource(new(fakePanicSource))
+
+		returnedLabels, err := getFeatureLabels(fakePanicFeatureSource)
+		Convey("No label is returned", func() {
+			So(len(returnedLabels), ShouldEqual, 0)
+		})
+		Convey("Error is produced and panic error is returned", func() {
+			So(err, ShouldResemble, fmt.Errorf("fake panic error"))
+		})
+
+	})
+}
diff --git a/sources.go b/sources.go
index 7fe3d40cc..7ccd330a9 100644
--- a/sources.go
+++ b/sources.go
@@ -102,19 +102,3 @@ func (s pstateSource) Discover() ([]string, error) {
 
 	return features, nil
 }
-
-////////////////////////////////////////////////////////////////////////////////
-// Fake Source (used only for testing)
-
-// Implements main.FeatureSource.
-type fakeSource struct{}
-
-func (s fakeSource) Name() string { return "fake" }
-func (s fakeSource) Discover() ([]string, error) {
-	features := []string{}
-
-	// Adding three fake features.
-	features = append(features, "fakefeature1", "fakefeature2", "fakefeature3")
-
-	return features, nil
-}
diff --git a/test_sources.go b/test_sources.go
new file mode 100644
index 000000000..645e34e43
--- /dev/null
+++ b/test_sources.go
@@ -0,0 +1,28 @@
+package main
+
+////////////////////////////////////////////////////////////////////////////////
+// Fake Source (used only for testing)
+
+// Implements main.FeatureSource.
+type fakeSource struct{}
+
+func (s fakeSource) Name() string { return "fake" }
+func (s fakeSource) Discover() ([]string, error) {
+	features := []string{}
+
+	// Adding three fake features.
+	features = append(features, "fakefeature1", "fakefeature2", "fakefeature3")
+
+	return features, nil
+}
+
+////////////////////////////////////////////////////////////////////////////////
+// Fake Panic Source (used only for testing)
+
+// Implements main.FeatureSource.
+type fakePanicSource struct{}
+
+func (s fakePanicSource) Name() string { return "fakepanic" }
+func (s fakePanicSource) Discover() ([]string, error) {
+	panic("fake panic error")
+}