From e4e39d820c46348069111f9744d053fe4c5c44eb Mon Sep 17 00:00:00 2001
From: James Shubin <james@shubin.ca>
Date: Mon, 13 Mar 2017 07:27:54 -0400
Subject: [PATCH] pgraph: semaphore: Refactor semaphore size function and test

---
 pgraph/semaphore.go      | 25 +++++++++++++++----------
 pgraph/semaphore_test.go | 13 +++++++++++++
 2 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/pgraph/semaphore.go b/pgraph/semaphore.go
index 08185c73..e2b6bc08 100644
--- a/pgraph/semaphore.go
+++ b/pgraph/semaphore.go
@@ -36,19 +36,10 @@ func (g *Graph) SemaLock(semas []string) error {
 	var reterr error
 	sort.Strings(semas) // very important to avoid deadlock in the dag!
 	for _, id := range semas {
-
-		size := 1 // default semaphore size
-		// valid id's include "some_id", "hello:42" and ":13"
-		if index := strings.LastIndex(id, SemaSep); index > -1 && (len(id)-index+len(SemaSep)) >= 1 {
-			// NOTE: we only allow size > 0 here!
-			if i, err := strconv.Atoi(id[index+len(SemaSep):]); err == nil && i > 0 {
-				size = i
-			}
-		}
-
 		g.slock.Lock()          // semaphore creation lock
 		sema, ok := g.semas[id] // lookup
 		if !ok {
+			size := SemaSize(id) // defaults to 1
 			g.semas[id] = semaphore.NewSemaphore(size)
 			sema = g.semas[id]
 		}
@@ -78,3 +69,17 @@ func (g *Graph) SemaUnlock(semas []string) error {
 	}
 	return reterr
 }
+
+// SemaSize returns the size integer associated with the semaphore id. It
+// defaults to 1 if not found.
+func SemaSize(id string) int {
+	size := 1 // default semaphore size
+	// valid id's include "some_id", "hello:42" and ":13"
+	if index := strings.LastIndex(id, SemaSep); index > -1 && (len(id)-index+len(SemaSep)) >= 1 {
+		// NOTE: we only allow size > 0 here!
+		if i, err := strconv.Atoi(id[index+len(SemaSep):]); err == nil && i > 0 {
+			size = i
+		}
+	}
+	return size
+}
diff --git a/pgraph/semaphore_test.go b/pgraph/semaphore_test.go
index 5212ad9b..940fd44e 100644
--- a/pgraph/semaphore_test.go
+++ b/pgraph/semaphore_test.go
@@ -23,6 +23,19 @@ import (
 	"github.com/purpleidea/mgmt/resources"
 )
 
+func TestSemaSize(t *testing.T) {
+	pairs := map[string]int{
+		"id:42":   42,
+		":13":     13,
+		"some_id": 1,
+	}
+	for id, size := range pairs {
+		if i := SemaSize(id); i != size {
+			t.Errorf("sema id `%s`, expected: `%d`, got: `%d`", id, size, i)
+		}
+	}
+}
+
 func NewNoopResTestSema(name string, semas []string) *NoopResTest {
 	obj := &NoopResTest{
 		NoopRes: resources.NoopRes{