diff --git a/pkg/blockstorage/awsebs/zone.go b/pkg/blockstorage/awsebs/zone.go index 6a34262dfe..011d5bf02a 100644 --- a/pkg/blockstorage/awsebs/zone.go +++ b/pkg/blockstorage/awsebs/zone.go @@ -2,26 +2,19 @@ package awsebs import ( "context" - "hash/fnv" - "sort" - "strings" - "github.com/pkg/errors" + "github.com/kanisterio/kanister/pkg/blockstorage/zone" log "github.com/sirupsen/logrus" - metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" - "k8s.io/client-go/kubernetes" - - "github.com/kanisterio/kanister/pkg/kube" ) func zoneForVolumeCreateFromSnapshot(ctx context.Context, region string, sourceZone string) (string, error) { - nzs, err := nodeZones(ctx) + nzs, err := zone.NodeZones(ctx) if err != nil { log.Errorf("Ignoring error getting Node availability zones. Error: %+v", err) } if len(nzs) != 0 { var z string - if z, err = zoneFromKnownNodeZones(ctx, region, sourceZone, nzs); err == nil && isZoneValid(z, region) { + if z, err = zone.GetZoneFromKnownNodeZones(ctx, sourceZone, nzs); err == nil && isZoneValid(z, region) { return z, nil } } @@ -39,40 +32,6 @@ func isZoneValid(zone, region string) bool { return false } -func zoneFromKnownNodeZones(ctx context.Context, region string, sourceZone string, nzs map[string]struct{}) (string, error) { - // If the original zone is available, we return that one. - if _, ok := nzs[sourceZone]; ok { - return sourceZone, nil - } - // If there's an available zone with the zone suffix, we use that one. - for nz := range nzs { - if zoneSuffixesMatch(nz, sourceZone) { - return nz, nil - } - } - // If any nodes are available, return an arbitrary one. - return consistentZone(sourceZone, nzs) -} - -func consistentZone(sourceZone string, nzs map[string]struct{}) (string, error) { - if len(nzs) == 0 { - return "", errors.New("could not restore volume: no zone found") - } - s := make([]string, 0, len(nzs)) - for nz := range nzs { - s = append(s, nz) - } - sort.Slice(s, func(i, j int) bool { - return strings.Compare(s[i], s[j]) < 0 - }) - h := fnv.New32() - if _, err := h.Write([]byte(sourceZone)); err != nil { - return "", errors.Errorf("failed to hash source zone %s: %s", sourceZone, err.Error()) - } - i := int(h.Sum32()) % len(nzs) - return s[i], nil -} - func zoneWithUnknownNodeZones(ctx context.Context, region string, sourceZone string) (string, error) { // We could not the zones of the nodes, so we return an arbitrary one. zs, err := regionToZones(ctx, region) @@ -83,40 +42,10 @@ func zoneWithUnknownNodeZones(ctx context.Context, region string, sourceZone str } // We look for a zone with the same suffix. for _, z := range zs { - if zoneSuffixesMatch(z, sourceZone) { + if zone.GetZoneSuffixesMatch(z, sourceZone) { return z, nil } } // We return an arbitrary zone in the region. return zs[0], nil } - -func zoneSuffixesMatch(zone1, zone2 string) bool { - return zone1[len(zone1)-1] == zone2[len(zone2)-1] -} - -const ( - nodeZonesErr = `Failed to get Node availability zones.` -) - -func nodeZones(ctx context.Context) (map[string]struct{}, error) { - cfg, err := kube.LoadConfig() - if err != nil { - return nil, errors.Wrap(err, nodeZonesErr) - } - cli, err := kubernetes.NewForConfig(cfg) - if err != nil { - return nil, errors.Wrap(err, nodeZonesErr) - } - ns, err := cli.CoreV1().Nodes().List(metav1.ListOptions{}) - if err != nil { - return nil, errors.Wrap(err, nodeZonesErr) - } - zoneSet := make(map[string]struct{}, len(ns.Items)) - for _, n := range ns.Items { - if v, ok := n.Labels[kube.PVZoneLabelName]; ok { - zoneSet[v] = struct{}{} - } - } - return zoneSet, nil -} diff --git a/pkg/blockstorage/awsebs/zone_test.go b/pkg/blockstorage/awsebs/zone_test.go index d23dc483bc..74ab56effc 100644 --- a/pkg/blockstorage/awsebs/zone_test.go +++ b/pkg/blockstorage/awsebs/zone_test.go @@ -47,40 +47,3 @@ func (s ZoneSuite) TestZoneWithUnknownNodeZones(c *C) { } } } - -func (s ZoneSuite) TestConsistentZone(c *C) { - // We don't care what the answer is as long as it's consistent. - for _, tc := range []struct { - sourceZone string - nzs map[string]struct{} - out string - }{ - { - sourceZone: "", - nzs: map[string]struct{}{ - "zone1": struct{}{}, - }, - out: "zone1", - }, - { - sourceZone: "", - nzs: map[string]struct{}{ - "zone1": struct{}{}, - "zone2": struct{}{}, - }, - out: "zone2", - }, - { - sourceZone: "from1", - nzs: map[string]struct{}{ - "zone1": struct{}{}, - "zone2": struct{}{}, - }, - out: "zone1", - }, - } { - out, err := consistentZone(tc.sourceZone, tc.nzs) - c.Assert(err, IsNil) - c.Assert(out, Equals, tc.out) - } -} diff --git a/pkg/blockstorage/zone/zone.go b/pkg/blockstorage/zone/zone.go new file mode 100644 index 0000000000..21110efcbe --- /dev/null +++ b/pkg/blockstorage/zone/zone.go @@ -0,0 +1,82 @@ +package zone + +import ( + "context" + "hash/fnv" + "sort" + "strings" + + "github.com/kanisterio/kanister/pkg/kube" + "github.com/pkg/errors" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/client-go/kubernetes" +) + +// GetZoneFromKnownNodeZones get the zone from known node zones +func GetZoneFromKnownNodeZones(ctx context.Context, sourceZone string, nzs map[string]struct{}) (string, error) { + // If the original zone is available, we return that one. + if _, ok := nzs[sourceZone]; ok { + return sourceZone, nil + } + // If there's an available zone with the zone suffix, we use that one. + for nz := range nzs { + if GetZoneSuffixesMatch(nz, sourceZone) { + return nz, nil + } + } + // If any nodes are available, return an arbitrary one. + return consistentZone(sourceZone, nzs) +} + +func consistentZone(sourceZone string, nzs map[string]struct{}) (string, error) { + if len(nzs) == 0 { + return "", errors.New("could not restore volume: no zone found") + } + s := make([]string, 0, len(nzs)) + for nz := range nzs { + s = append(s, nz) + } + sort.Slice(s, func(i, j int) bool { + return strings.Compare(s[i], s[j]) < 0 + }) + h := fnv.New32() + if _, err := h.Write([]byte(sourceZone)); err != nil { + return "", errors.Errorf("failed to hash source zone %s: %s", sourceZone, err.Error()) + } + i := int(h.Sum32()) % len(nzs) + return s[i], nil +} + +// GetZoneSuffixesMatch check if the given zones have a matching suffix +func GetZoneSuffixesMatch(zone1, zone2 string) bool { + a1 := zone1[len(zone1)-1] + a2 := zone2[len(zone2)-1] + return a1 == a2 +} + +const ( + nodeZonesErr = `Failed to get Node availability zones.` +) + +// NodeZones get the zones available for the nodes +func NodeZones(ctx context.Context) (map[string]struct{}, error) { + cfg, err := kube.LoadConfig() + if err != nil { + return nil, errors.Wrap(err, nodeZonesErr) + } + cli, err := kubernetes.NewForConfig(cfg) + if err != nil { + return nil, errors.Wrap(err, nodeZonesErr) + } + ns, err := cli.CoreV1().Nodes().List(metav1.ListOptions{}) + if err != nil { + return nil, errors.Wrap(err, nodeZonesErr) + } + zoneSet := make(map[string]struct{}, len(ns.Items)) + for _, n := range ns.Items { + if v, ok := n.Labels[kube.PVZoneLabelName]; ok { + zoneSet[v] = struct{}{} + } + } + return zoneSet, nil +} diff --git a/pkg/blockstorage/awsebs/awsebs_kube_test.go b/pkg/blockstorage/zone/zone_kube_test.go similarity index 85% rename from pkg/blockstorage/awsebs/awsebs_kube_test.go rename to pkg/blockstorage/zone/zone_kube_test.go index 67af9ebcaa..784c09c687 100644 --- a/pkg/blockstorage/awsebs/awsebs_kube_test.go +++ b/pkg/blockstorage/zone/zone_kube_test.go @@ -1,6 +1,6 @@ // +build !unit -package awsebs +package zone import ( "context" @@ -14,7 +14,7 @@ var _ = Suite(&KubeTestAWSEBSSuite{}) func (s KubeTestAWSEBSSuite) TestNodeZones(c *C) { ctx := context.Background() - zones, err := nodeZones(ctx) + zones, err := NodeZones(ctx) c.Assert(err, IsNil) c.Assert(zones, Not(HasLen), 0) } diff --git a/pkg/blockstorage/zone/zone_test.go b/pkg/blockstorage/zone/zone_test.go new file mode 100644 index 0000000000..fb59c13118 --- /dev/null +++ b/pkg/blockstorage/zone/zone_test.go @@ -0,0 +1,51 @@ +package zone + +import ( + "testing" + + . "gopkg.in/check.v1" +) + +// Hook up gocheck into the "go test" runner. +func Test(t *testing.T) { TestingT(t) } + +type ZoneSuite struct{} + +var _ = Suite(&ZoneSuite{}) + +func (s ZoneSuite) TestConsistentZone(c *C) { + // We don't care what the answer is as long as it's consistent. + for _, tc := range []struct { + sourceZone string + nzs map[string]struct{} + out string + }{ + { + sourceZone: "", + nzs: map[string]struct{}{ + "zone1": {}, + }, + out: "zone1", + }, + { + sourceZone: "", + nzs: map[string]struct{}{ + "zone1": {}, + "zone2": {}, + }, + out: "zone2", + }, + { + sourceZone: "from1", + nzs: map[string]struct{}{ + "zone1": {}, + "zone2": {}, + }, + out: "zone1", + }, + } { + out, err := consistentZone(tc.sourceZone, tc.nzs) + c.Assert(err, IsNil) + c.Assert(out, Equals, tc.out) + } +}