Repository: threefoldtech/zos
Pull request #2015 (2023-07-31 06:34:24)

add more geoip service urls
## Related issues:

- #1887 


Patch:
diff --git a/go.mod b/go.mod
index cb6e27d7..2ca3b06d 100644
--- a/go.mod
+++ b/go.mod
@@ -42,7 +42,7 @@ require (
 	github.com/shirou/gopsutil v3.21.11+incompatible
 	github.com/stretchr/testify v1.8.2
 	github.com/threefoldtech/0-fs v1.3.1-0.20201203163303-d963de9adea7
-	github.com/threefoldtech/tfchain/clients/tfchain-client-go v0.0.0-20230726134044-983e3593f1e8
+	github.com/threefoldtech/tfchain/clients/tfchain-client-go v0.0.0-20230809064214-d8ef5ca360eb
 	github.com/threefoldtech/tfgrid-sdk-go/rmb-sdk-go v0.4.0
 	github.com/threefoldtech/zbus v1.0.1
 	github.com/tyler-smith/go-bip39 v1.1.0
diff --git a/go.sum b/go.sum
index d8f0c1d6..e0aee5c4 100644
--- a/go.sum
+++ b/go.sum
@@ -587,6 +587,8 @@ github.com/threefoldtech/tfchain/clients/tfchain-client-go v0.0.0-20230711081157
 github.com/threefoldtech/tfchain/clients/tfchain-client-go v0.0.0-20230711081157-7c84c6d78b4d/go.mod h1:dtDKAPiUDxAwIkfHV7xcAFZcOm+xwNIuOI1MLFS+MeQ=
 github.com/threefoldtech/tfchain/clients/tfchain-client-go v0.0.0-20230726134044-983e3593f1e8 h1:WOU12x4+k9ZueOs0r+06FsKdbi3DiF/pqPifMhglMng=
 github.com/threefoldtech/tfchain/clients/tfchain-client-go v0.0.0-20230726134044-983e3593f1e8/go.mod h1:dtDKAPiUDxAwIkfHV7xcAFZcOm+xwNIuOI1MLFS+MeQ=
+github.com/threefoldtech/tfchain/clients/tfchain-client-go v0.0.0-20230809064214-d8ef5ca360eb h1:ljvL1BMwTPrV0kAz9E56zJ5jq39TPRsd/tP6l0+0TzY=
+github.com/threefoldtech/tfchain/clients/tfchain-client-go v0.0.0-20230809064214-d8ef5ca360eb/go.mod h1:dtDKAPiUDxAwIkfHV7xcAFZcOm+xwNIuOI1MLFS+MeQ=
 github.com/threefoldtech/tfgrid-sdk-go/rmb-sdk-go v0.4.0 h1:fmCYshC+u0/1nqlu10TnZjZeywioXOSWpm6n8zPKgkI=
 github.com/threefoldtech/tfgrid-sdk-go/rmb-sdk-go v0.4.0/go.mod h1:tttnc/l4F2lx1SyETTiLrMaiQEgMeGCvlVEIW6nUeJA=
 github.com/threefoldtech/zbus v1.0.1 h1:3KaEpyOiDYAw+lrAyoQUGIvY9BcjVRXlQ1beBRqhRNk=
diff --git a/pkg/events/redis.go b/pkg/events/redis.go
index d08fa8c8..84a6d621 100644
--- a/pkg/events/redis.go
+++ b/pkg/events/redis.go
@@ -236,6 +236,7 @@ func (r *RedisConsumer) consumer(ctx context.Context, stream string, ch reflect.
 	logger := log.With().Str("stream", stream).Logger()
 	go func() {
 		defer con.Close()
+		defer ch.Close()
 
 		for {
 			messages, err := r.pop(con, group, stream)
diff --git a/pkg/geoip/geoip.go b/pkg/geoip/geoip.go
index 5f535763..1b9a5918 100644
--- a/pkg/geoip/geoip.go
+++ b/pkg/geoip/geoip.go
@@ -4,6 +4,8 @@ import (
 	"encoding/json"
 	"errors"
 	"net/http"
+
+	"github.com/rs/zerolog/log"
 )
 
 // Location holds the result of a geoip request
@@ -16,8 +18,27 @@ type Location struct {
 	City        string  `json:"city_name"`
 }
 
+var (
+	geoipURLs = []string{"https://geoip.grid.tf/", "https://02.geoip.grid.tf/", "https://03.geoip.grid.tf/"}
+)
+
 // Fetch retrieves the location of the system calling this function
 func Fetch() (Location, error) {
+
+	for _, url := range geoipURLs {
+		l, err := getLocation(url)
+		if err != nil {
+			log.Err(err).Str("url", url).Msg("failed to fetch location from geoip service")
+			continue
+		}
+
+		return l, nil
+	}
+
+	return Location{}, errors.New("failed to fetch location information")
+}
+
+func getLocation(geoIPService string) (Location, error) {
 	l := Location{
 		Longitute: 0.0,
 		Latitude:  0.0,
@@ -26,19 +47,19 @@ func Fetch() (Location, error) {
 		City:      "Unknown",
 	}
 
-	resp, err := http.Get("https://geoip.grid.tf")
+	resp, err := http.Get(geoIPService)
 	if err != nil {
 		return l, err
 	}
 	defer resp.Body.Close()
 
 	if resp.StatusCode != 200 {
-		return l, errors.New("error fetch location")
+		return l, errors.New("error fetching location")
 	}
 
 	if err := json.NewDecoder(resp.Body).Decode(&l); err != nil {
 		return l, err
 	}
 
-	return l, err
+	return l, nil
 }
diff --git a/pkg/power/power.go b/pkg/power/power.go
index 3c40b283..834ef8e7 100644
--- a/pkg/power/power.go
+++ b/pkg/power/power.go
@@ -126,8 +126,7 @@ func (p *PowerServer) syncSelf() error {
 	// if target is down, we make sure state is down, then shutdown
 
 	if power.Target.IsUp {
-		if power.State.IsDown {
-			_, err = cl.SetNodePowerState(p.identity, true)
+		if err := p.setNodePowerState(cl, true); err != nil {
 			return errors.Wrap(err, "failed to set state to up")
 		}
 
@@ -137,9 +136,8 @@ func (p *PowerServer) syncSelf() error {
 	// now the target must be down.
 	// we need to shutdown
 
-	if power.State.IsUp {
-		_, err = cl.SetNodePowerState(p.identity, false)
-		return errors.Wrap(err, "failed to set state to up")
+	if err := p.setNodePowerState(cl, false); err != nil {
+		return errors.Wrap(err, "failed to set state to down")
 	}
 
 	// otherwise node need to get back to sleep.
@@ -220,7 +218,7 @@ func (p *PowerServer) event(event *pkg.PowerTargetChangeEvent) error {
 
 	if event.NodeID == p.node && event.Target.IsDown {
 		// we need to shutdown!
-		if _, err := cl.SetNodePowerState(p.identity, false); err != nil {
+		if err := p.setNodePowerState(cl, false); err != nil {
 			return errors.Wrap(err, "failed to set node power state to down")
 		}
 
@@ -237,6 +235,43 @@ func (p *PowerServer) event(event *pkg.PowerTargetChangeEvent) error {
 	return nil
 }
 
+// setNodePowerState sets the node power state as provided or to up if power mgmt is
+// not enabled on this node.
+// this function makes sure to compare the state with on chain state to not do
+// un-necessary transactions.
+func (p *PowerServer) setNodePowerState(cl *substrate.Substrate, up bool) error {
+
+	/*
+		if power is not enabled, the node state should always be up
+		otherwise update the state to the correct value
+
+		| enabled | up | result|
+		| 0       | 0  | 1     |
+		| 0       | 1  | 1     |
+		| 1       | 0  | 0     |
+		| 1       | 1  | 1     |
+
+		this simplifies as below:
+	*/
+
+	up = !p.enabled || up
+	power, err := cl.GetPowerTarget(p.node)
+
+	if err != nil {
+		return errors.Wrap(err, "failed to check power state")
+	}
+
+	// only update the chain if it's different from actual value.
+	if power.State.IsUp == up {
+		return nil
+	}
+
+	log.Info().Bool("state", up).Msg("setting node power state")
+	// this to make sure node state is fixed also for nodes
+	_, err = cl.SetNodePowerState(p.identity, up)
+	return err
+}
+
 func (p *PowerServer) recv(ctx context.Context) error {
 	log.Info().Msg("listening for power events")
 	stream, err := p.consumer.PowerTargetChange(ctx)
diff --git a/pkg/provision/engine.go b/pkg/provision/engine.go
index 51dd8c20..cb09c2f5 100644
--- a/pkg/provision/engine.go
+++ b/pkg/provision/engine.go
@@ -15,6 +15,7 @@ import (
 	substrate "github.com/threefoldtech/tfchain/clients/tfchain-client-go"
 	"github.com/threefoldtech/zos/pkg"
 	"github.com/threefoldtech/zos/pkg/gridtypes"
+	"github.com/threefoldtech/zos/pkg/gridtypes/zos"
 )
 
 // EngineOption interface
@@ -867,10 +868,42 @@ func (e *NativeEngine) uninstallDeployment(ctx context.Context, dl *gridtypes.De
 
 }
 
+func getZmountSize(wl *gridtypes.Workload) (gridtypes.Unit, error) {
+	data, err := wl.WorkloadData()
+	if err != nil {
+		return 0, err
+	}
+	zmount, ok := data.(*zos.ZMount)
+	if !ok {
+		return 0, fmt.Errorf("failed to get workload data as a zmount '%v'", data)
+	}
+	return zmount.Size, nil
+}
+
+func sortZmountWorkloads(workloads []*gridtypes.WorkloadWithID) {
+	sort.Slice(workloads, func(i, j int) bool {
+		sizeI, err := getZmountSize(workloads[i].Workload)
+		if err != nil {
+			return false
+		}
+
+		sizeJ, err := getZmountSize(workloads[j].Workload)
+		if err != nil {
+			return false
+		}
+
+		return sizeI > sizeJ
+	})
+}
+
 func (e *NativeEngine) installDeployment(ctx context.Context, getter gridtypes.WorkloadGetter) {
 	for _, typ := range e.order {
 		workloads := getter.ByType(typ)
 
+		if typ == zos.ZMountType {
+			sortZmountWorkloads(workloads)
+		}
+
 		for _, wl := range workloads {
 			if err := e.installWorkload(ctx, wl); err != nil {
 				log.Error().Err(err).Stringer("id", wl.ID).Msg("failed to install workload")


Test patch:
diff --git a/pkg/provision/engine_test.go b/pkg/provision/engine_test.go
index e9f2765c..177f6bf0 100644
--- a/pkg/provision/engine_test.go
+++ b/pkg/provision/engine_test.go
@@ -1,5 +1,14 @@
 package provision
 
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/threefoldtech/zos/pkg/gridtypes"
+	"github.com/threefoldtech/zos/pkg/gridtypes/zos"
+)
+
 // func TestEngine(t *testing.T) {
 // 	td, err := ioutil.TempDir("", "")
 // 	require.NoError(t, err)
@@ -68,3 +77,64 @@ package provision
 // 	assert.EqualValues(t, 1, workloads.ZDBNamespace)
 // 	assert.EqualValues(t, 0, workloads.K8sVM)
 // }
+
+func TestGetZmountSize(t *testing.T) {
+	t.Run("invalid type", func(t *testing.T) {
+		wl := gridtypes.WorkloadWithID{
+			Workload: &gridtypes.Workload{Type: "invalid"},
+		}
+		_, err := getZmountSize(wl.Workload)
+		assert.Error(t, err)
+	})
+	t.Run("different data type", func(t *testing.T) {
+		wl := gridtypes.WorkloadWithID{
+			Workload: &gridtypes.Workload{Type: zos.ZDBType, Data: json.RawMessage(`{"size": 10}`)},
+		}
+		_, err := getZmountSize(wl.Workload)
+		assert.Error(t, err)
+	})
+	t.Run("valid data", func(t *testing.T) {
+		wl := gridtypes.WorkloadWithID{
+			Workload: &gridtypes.Workload{Type: zos.ZMountType, Data: json.RawMessage(`{"size": 10}`)},
+		}
+		size, err := getZmountSize(wl.Workload)
+		assert.NoError(t, err)
+		assert.Equal(t, size, gridtypes.Unit(10))
+	})
+
+}
+
+func TestSortZmountWorkloads(t *testing.T) {
+	workloads := []*gridtypes.WorkloadWithID{
+		{Workload: &gridtypes.Workload{
+			Type: zos.ZMountType,
+			Data: json.RawMessage(`{"size": 10}`),
+		}},
+		{Workload: &gridtypes.Workload{
+			Type: zos.ZMountType,
+			Data: json.RawMessage(`{"size": 30}`),
+		}},
+		{Workload: &gridtypes.Workload{
+			Type: zos.ZMountType,
+			Data: json.RawMessage(`{"size": 20}`),
+		}},
+	}
+
+	expectedWorkloads := []*gridtypes.WorkloadWithID{
+		{Workload: &gridtypes.Workload{
+			Type: zos.ZMountType,
+			Data: json.RawMessage(`{"size": 30}`),
+		}},
+		{Workload: &gridtypes.Workload{
+			Type: zos.ZMountType,
+			Data: json.RawMessage(`{"size": 20}`),
+		}},
+		{Workload: &gridtypes.Workload{
+			Type: zos.ZMountType,
+			Data: json.RawMessage(`{"size": 10}`),
+		}},
+	}
+
+	sortZmountWorkloads(workloads)
+	assert.Equal(t, expectedWorkloads, workloads)
+}