diff --git a/solrman/smstorage/inmemory_storage.go b/solrman/smstorage/inmemory_storage.go index 8aae48c..4e88257 100644 --- a/solrman/smstorage/inmemory_storage.go +++ b/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) { diff --git a/solrman/smstorage/inmemory_storage_test.go b/solrman/smstorage/inmemory_storage_test.go index 0204670..97d54ad 100644 --- a/solrman/smstorage/inmemory_storage_test.go +++ b/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{}) +} diff --git a/solrman/smstorage/storage.go b/solrman/smstorage/storage.go index e284e2e..267d4d9 100644 --- a/solrman/smstorage/storage.go +++ b/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) // diff --git a/solrman/smstorage/storage_test.go b/solrman/smstorage/storage_test.go index 4b37a38..04a3891 100644 --- a/solrman/smstorage/storage_test.go +++ b/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 + } +} diff --git a/solrman/smstorage/zk_storage.go b/solrman/smstorage/zk_storage.go index 5e6bfb0..e37c81b 100644 --- a/solrman/smstorage/zk_storage.go +++ b/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) diff --git a/solrman/smstorage/zk_storage_test.go b/solrman/smstorage/zk_storage_test.go index 7cc1a92..5a5659d 100644 --- a/solrman/smstorage/zk_storage_test.go +++ b/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) +}