add / remove from StationaryOrgsList

solrman/smstorage/inmemory_storage.go

@@ -33,6 +33,7 @@ type InMemoryStorage struct {
 	queryAggregatorStabbingEnabled bool
 	inProgress                     map[string]solrmanapi.OpRecord
 	completed                      []solrmanapi.OpRecord
+	stationaryOrgList              []string
 }
 
 var _ SolrManStorage = &InMemoryStorage{}
@@ -98,7 +99,30 @@ func (s *InMemoryStorage) GetStationaryOrgList() ([]string, error) {
 	s.mu.RLock()
 	defer s.mu.RUnlock()
 
-	return nil, nil
+	return s.stationaryOrgList, nil
+}
+
+func (s *InMemoryStorage) AddStationaryOrgs(orgs []string) ([]string, error) {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	s.stationaryOrgList = append(s.stationaryOrgList, orgs...)
+
+	return s.stationaryOrgList, nil
+}
+
+func (s *InMemoryStorage) RemoveStationaryOrgs(orgs []string) ([]string, error) {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
+	for _, org := range orgs {
+		for i := range s.stationaryOrgList {
+			if s.stationaryOrgList[i] == org {
+				s.stationaryOrgList = append(s.stationaryOrgList[:i], s.stationaryOrgList[i+1:]...)
+				break
+			}
+		}
+	}
+	return s.stationaryOrgList, nil
 }
 
 func (s *InMemoryStorage) IsDisabled() (bool, error) {

solrman/smstorage/inmemory_storage_test.go

@@ -25,3 +25,7 @@ func TestInMemoryStorage_InProgressOps(t *testing.T) {
 func TestInMemoryStorage_CompletedOps(t *testing.T) {
 	testStorage_CompletedOps(t, &InMemoryStorage{})
 }
+
+func TestInMemoryStorage_StationaryOrgList(t *testing.T) {
+	testStorage_StationaryOrgList(t, &InMemoryStorage{})
+}

solrman/smstorage/storage.go

@@ -43,6 +43,8 @@ type SolrManStorage interface {
 
 	// Returns a list of solr orgs that should not be split or moved.
 	GetStationaryOrgList() ([]string, error)
+	AddStationaryOrgs(orgs []string) ([]string, error)
+	RemoveStationaryOrgs(orgs []string) ([]string, error)
 
 	IsDisabled() (bool, error)                      // if true, solrman is entirely disabled
 	GetDisabledReasons() (map[string]string, error) //

solrman/smstorage/storage_test.go

@@ -127,3 +127,67 @@ func testStorage_CompletedOps(t *testing.T, s SolrManStorage) {
 	// A negative count also returns all values.
 	assertOps(-1, expectOps...)
 }
+
+func testStorage_StationaryOrgList(t *testing.T, s SolrManStorage) {
+	stationaryOrgList, err := s.GetStationaryOrgList()
+	if err != nil {
+		t.Errorf("GetStationaryOrgList failed: %s", err)
+		return
+	}
+	if len(stationaryOrgList) != 0 {
+		t.Errorf("expected empty stationary org list, got %v", stationaryOrgList)
+	}
+
+	// test adding orgs
+	orgs := []string{"org1", "org2", "org3"}
+	stationaryOrgList, err = s.AddStationaryOrgs(orgs)
+	if err != nil {
+		t.Errorf("AddStationaryOrgs failed: %s", err)
+		return
+	}
+	if len(orgs) != len(stationaryOrgList) {
+		t.Errorf("expected %d orgs in stationary org list, got %v", len(orgs), stationaryOrgList)
+	}
+	for i := range orgs {
+		for j := range stationaryOrgList {
+			if orgs[i] == stationaryOrgList[j] {
+				break
+			}
+			if j == len(stationaryOrgList)-1 {
+				t.Errorf("expected org %s in stationary org list, got %v", orgs[i], stationaryOrgList)
+			}
+		}
+	}
+
+	orgsToRemove := []string{"org2", "org3"}
+	expectedRemainingOrgs := []string{"org1"}
+	stationaryOrgList, err = s.RemoveStationaryOrgs(orgsToRemove)
+	if err != nil {
+		t.Errorf("RemoveStationaryOrgs failed: %s", err)
+		return
+	}
+	if len(expectedRemainingOrgs) != len(stationaryOrgList) {
+		t.Errorf("expected %d orgs in stationary org list, got %v", len(orgs), stationaryOrgList)
+	}
+	for _, org := range stationaryOrgList {
+		if org == "org2" || org == "org3" {
+			t.Errorf("expected org %s not in stationary org list, got %v", org, stationaryOrgList)
+		}
+	}
+	for i := range expectedRemainingOrgs {
+		for j := range stationaryOrgList {
+			if expectedRemainingOrgs[i] == stationaryOrgList[j] {
+				break
+			}
+			if j == len(stationaryOrgList)-1 {
+				t.Errorf("expected org %s in stationary org list, got %v", orgs[i], stationaryOrgList)
+			}
+		}
+	}
+
+	stationaryOrgList, err = s.RemoveStationaryOrgs([]string{"doesnotexist"})
+	if err != nil {
+		t.Errorf("RemoveStationaryOrgs failed: %s", err)
+		return
+	}
+}

solrman/smstorage/zk_storage.go

@@ -206,6 +206,61 @@ func (s *ZkStorage) GetStationaryOrgList() ([]string, error) {
 	return children, nil
 }
 
+func (s *ZkStorage) AddStationaryOrgs(orgs []string) ([]string, error) {
+	path := s.stationaryPath()
+	children, _, err := s.conn.Children(path)
+	if err == zk.ErrNoNode {
+		return nil, nil
+	}
+	if err != nil {
+		return nil, smutil.Cherrf(err, "could not get children at %s in ZK", path)
+	}
+
+	for _, org := range orgs {
+		orgPath := fmt.Sprintf("%s/%s", path, org)
+		_, err = s.conn.Create(orgPath, nil, 0, zk.WorldACL(zk.PermAll))
+		if err == zk.ErrNodeExists {
+			continue
+		}
+		if err != nil {
+			return nil, smutil.Cherrf(err, "could not create %s in ZK", orgPath)
+		}
+	}
+
+	children, _, err = s.conn.Children(path)
+	if err == zk.ErrNoNode {
+		return nil, nil
+	}
+	if err != nil {
+		return nil, smutil.Cherrf(err, "could not get children at %s in ZK", path)
+	}
+	sort.Strings(children)
+
+	return children, nil
+}
+
+func (s *ZkStorage) RemoveStationaryOrgs(orgs []string) ([]string, error) {
+	path := s.stationaryPath()
+	for _, org := range orgs {
+		orgPath := fmt.Sprintf("%s/%s", path, org)
+		err := s.conn.Delete(orgPath, -1)
+		if err != nil && err != zk.ErrNoNode {
+			return nil, smutil.Cherrf(err, "could not delete %s in ZK", orgPath)
+		}
+	}
+
+	children, _, err := s.conn.Children(path)
+	if err == zk.ErrNoNode {
+		return nil, nil
+	}
+	if err != nil {
+		return nil, smutil.Cherrf(err, "could not get children at %s in ZK", path)
+	}
+	sort.Strings(children)
+
+	return children, nil
+}
+
 func (s *ZkStorage) IsDisabled() (bool, error) {
 	path := s.disabledPath()
 	children, _, err := s.conn.Children(path)

solrman/smstorage/zk_storage_test.go

@@ -138,10 +138,10 @@ func TestZkStorage_IsDisabled(t *testing.T) {
 
 	testutil.createWithData(s.disabledPath()+"/test", "testDisabled")
 
-	if isDisabled, _ := s.IsDisabled(); !IsDisabled {
+	if isDisabled, _ := s.IsDisabled(); !isDisabled {
 		t.Error("expected to be disabled")
-	} else if reasons, _ := s.GetDisabledReasons(); reasons[0] != "testDisabled" {
-		t.Errorf("expect reason is \"testDisabled\"; got %s", reason)
+	} else if reasons, _ := s.GetDisabledReasons(); reasons["test"] != "testDisabled" {
+		t.Errorf("expect reason is \"testDisabled\"; got %s", reasons["test"])
 	}
 
 }
@@ -151,7 +151,7 @@ func TestZkStorage_SetDisabled(t *testing.T) {
 	defer testutil.teardown()
 
 	if isDisabled, reason := s.IsDisabled(); isDisabled {
-		t.Error("expected to not be disabled; reason found was %s", reason)
+		t.Errorf("expected to not be disabled; reason found was %s", reason)
 	}
 	if ok, _, _ := s.conn.Exists(s.disabledPath()); ok {
 		t.Errorf("%s should not exist", s.disabledPath())
@@ -233,3 +233,9 @@ func TestZkStorage_IsQueryAggregatorStabbingEnabled(t *testing.T) {
 		t.Error("expected query aggregator stabbing to be enabled")
 	}
 }
+
+func TestZkStorage_GetStationaryOrgList(t *testing.T) {
+	s, testutil := setup(t)
+	defer testutil.teardown()
+	testStorage_StationaryOrgList(t, s)
+}