Merge pull request #179 from gustavosbarreto/probe_try_again_after_ti…

…meout

probe state: try again after timeout

client/update.go

@@ -60,7 +60,7 @@ func (u *UpdateClient) ProbeUpdate(api ApiRequester, uri string, data interface{
 
 	res, err := api.Do(req)
 	if err != nil {
-		finalErr := fmt.Errorf("probe update request failed: %s", err)
+		finalErr := errors.Wrap(err, "probe update request failed")
 		log.Error(finalErr)
 		return nil, nil, 0, finalErr
 	}

server/agent_backend_test.go

@@ -223,7 +223,7 @@ func TestProbeRouteWithDefaultApiClient(t *testing.T) {
 	cm := &controllermock.ControllerMock{}
 
 	uh.Controller = cm
-	cm.On("ProbeUpdate", uh.DefaultApiClient, 5).Return((*metadata.UpdateMetadata)(nil), []byte{}, 3600*time.Second)
+	cm.On("ProbeUpdate", uh.DefaultApiClient, 5).Return((*metadata.UpdateMetadata)(nil), []byte{}, 3600*time.Second, nil)
 
 	rwm := &responsewritermock.ResponseWriterMock{}
 	rwm.On("WriteHeader", 200)
@@ -320,7 +320,7 @@ func TestProbeRouteWithServerAddressField(t *testing.T) {
 				// Not needed for this test case
 				actual.CheckRedirect = nil
 				return reflect.DeepEqual(actual, apiClient)
-			}), 5).Return((*metadata.UpdateMetadata)(nil), []byte{}, 3600*time.Second)
+			}), 5).Return((*metadata.UpdateMetadata)(nil), []byte{}, 3600*time.Second, nil)
 
 			rwm := &responsewritermock.ResponseWriterMock{}
 			rwm.On("WriteHeader", 200)

testsmocks/controllermock/controller.go

@@ -20,9 +20,9 @@ type ControllerMock struct {
 	mock.Mock
 }
 
-func (cm *ControllerMock) ProbeUpdate(apiClient *client.ApiClient, retries int) (*metadata.UpdateMetadata, []byte, time.Duration) {
+func (cm *ControllerMock) ProbeUpdate(apiClient *client.ApiClient, retries int) (*metadata.UpdateMetadata, []byte, time.Duration, error) {
 	args := cm.Called(apiClient, retries)
-	return args.Get(0).(*metadata.UpdateMetadata), args.Get(1).([]byte), args.Get(2).(time.Duration)
+	return args.Get(0).(*metadata.UpdateMetadata), args.Get(1).([]byte), args.Get(2).(time.Duration), args.Error(3)
 }
 
 func (cm *ControllerMock) DownloadUpdate(apiClient *client.ApiClient, updateMetadata *metadata.UpdateMetadata, cancel <-chan bool, progressChan chan<- int) error {

testsmocks/controllermock/controller_test.go

@@ -9,6 +9,7 @@
 package controllermock
 
 import (
+	"errors"
 	"fmt"
 	"testing"
 	"time"
@@ -22,16 +23,18 @@ func TestProbeUpdate(t *testing.T) {
 	expectedMetadata := &metadata.UpdateMetadata{}
 	expectedSignature := []byte{}
 	expectedDuration := 10 * time.Second
+	expectedErr := errors.New("")
 
 	apiClient := client.NewApiClient("address")
 	cm := &ControllerMock{}
-	cm.On("ProbeUpdate", apiClient, 0).Return(expectedMetadata, expectedSignature, expectedDuration)
+	cm.On("ProbeUpdate", apiClient, 0).Return(expectedMetadata, expectedSignature, expectedDuration, expectedErr)
 
-	m, s, d := cm.ProbeUpdate(apiClient, 0)
+	m, s, d, err := cm.ProbeUpdate(apiClient, 0)
 
 	assert.Equal(t, expectedMetadata, m)
 	assert.Equal(t, expectedSignature, s)
 	assert.Equal(t, expectedDuration, d)
+	assert.Equal(t, expectedErr, err)
 
 	cm.AssertExpectations(t)
 }

updatehub/poll_state_test.go

@@ -80,9 +80,9 @@ func TestPollingRetries(t *testing.T) {
 		})
 	}().Unpatch()
 
-	cm.On("ProbeUpdate", apiClient, 0).Return((*metadata.UpdateMetadata)(nil), []byte{}, time.Duration(-1)).Once()
-	cm.On("ProbeUpdate", apiClient, 1).Return((*metadata.UpdateMetadata)(nil), []byte{}, time.Duration(-1)).Once()
-	cm.On("ProbeUpdate", apiClient, 2).Return((*metadata.UpdateMetadata)(nil), []byte{}, time.Duration(-1)).Once()
+	cm.On("ProbeUpdate", apiClient, 0).Return((*metadata.UpdateMetadata)(nil), []byte{}, time.Duration(-1), nil).Once()
+	cm.On("ProbeUpdate", apiClient, 1).Return((*metadata.UpdateMetadata)(nil), []byte{}, time.Duration(-1), nil).Once()
+	cm.On("ProbeUpdate", apiClient, 2).Return((*metadata.UpdateMetadata)(nil), []byte{}, time.Duration(-1), nil).Once()
 
 	uh.Controller = cm
 	uh.Settings.PollingInterval = time.Second
@@ -109,7 +109,7 @@ func TestPollingRetries(t *testing.T) {
 	sha256sum := sha256.Sum256([]byte(validJSONMetadata))
 	signature, _ := rsa.SignPKCS1v15(rand.Reader, testPrivateKey, crypto.SHA256, sha256sum[:])
 
-	cm.On("ProbeUpdate", apiClient, 3).Return(um, signature, time.Duration(0)).Once()
+	cm.On("ProbeUpdate", apiClient, 3).Return(um, signature, time.Duration(0), nil).Once()
 
 	state, _ := next.Handle(uh)
 	assert.IsType(t, &IdleState{}, state)

updatehub/probe_state.go

@@ -9,9 +9,11 @@
 package updatehub
 
 import (
-	"errors"
+	"net"
 	"time"
 
+	"github.com/OSSystems/pkg/log"
+	"github.com/pkg/errors"
 	"github.com/updatehub/updatehub/client"
 	"github.com/updatehub/updatehub/metadata"
 )
@@ -39,8 +41,19 @@ func (state *ProbeState) ID() UpdateHubState {
 // polling state otherwise.
 func (state *ProbeState) Handle(uh *UpdateHub) (State, bool) {
 	var signature []byte
+	var err error
 
-	state.probeUpdateMetadata, signature, state.probeExtraPoll = uh.Controller.ProbeUpdate(state.apiClient, uh.Settings.PollingRetries)
+	for {
+		state.probeUpdateMetadata, signature, state.probeExtraPoll, err = uh.Controller.ProbeUpdate(state.apiClient, uh.Settings.PollingRetries)
+
+		if neterr, ok := errors.Cause(err).(net.Error); ok && neterr.Timeout() {
+			log.Warn("timeout during download update")
+			uh.Settings.PollingRetries++
+			continue
+		}
+
+		break
+	}
 
 	// "non-blocking" write to channel
 	select {

updatehub/probe_state_test.go

@@ -17,8 +17,10 @@ import (
 	"testing"
 	"time"
 
+	errs "github.com/pkg/errors"
 	"github.com/spf13/afero"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
 	"github.com/updatehub/updatehub/client"
 	"github.com/updatehub/updatehub/installmodes"
 	"github.com/updatehub/updatehub/metadata"
@@ -54,7 +56,7 @@ func TestStateProbeWithUpdateAvailable(t *testing.T) {
 	sha256sum := sha256.Sum256([]byte(validJSONMetadata))
 	signature, _ := rsa.SignPKCS1v15(rand.Reader, testPrivateKey, crypto.SHA256, sha256sum[:])
 
-	cm.On("ProbeUpdate", apiClient, 0).Return(um, signature, time.Duration(0))
+	cm.On("ProbeUpdate", apiClient, 0).Return(um, signature, time.Duration(0), nil)
 
 	next, _ := uh.GetState().Handle(uh)
 
@@ -87,7 +89,7 @@ func TestStateProbeWithUpdateNotAvailable(t *testing.T) {
 
 	uh.Store.Remove(uh.Settings.RuntimeSettingsPath)
 
-	cm.On("ProbeUpdate", apiClient, 0).Return((*metadata.UpdateMetadata)(nil), []byte{}, time.Duration(0))
+	cm.On("ProbeUpdate", apiClient, 0).Return((*metadata.UpdateMetadata)(nil), []byte{}, time.Duration(0), nil)
 
 	next, _ := uh.GetState().Handle(uh)
 
@@ -127,7 +129,7 @@ func TestStateProbeWithExtraPoll(t *testing.T) {
 
 	uh.Store.Remove(uh.Settings.RuntimeSettingsPath)
 
-	cm.On("ProbeUpdate", apiClient, 0).Return((*metadata.UpdateMetadata)(nil), []byte{}, time.Duration(5*time.Second))
+	cm.On("ProbeUpdate", apiClient, 0).Return((*metadata.UpdateMetadata)(nil), []byte{}, time.Duration(5*time.Second), nil)
 
 	next, _ := uh.GetState().Handle(uh)
 
@@ -172,7 +174,7 @@ func TestStateProbeWithUpdateAvailableButAlreadyInstalled(t *testing.T) {
 
 	uh.lastInstalledPackageUID = m.PackageUID()
 
-	cm.On("ProbeUpdate", apiClient, 0).Return(m, []byte{}, time.Duration(0))
+	cm.On("ProbeUpdate", apiClient, 0).Return(m, []byte{}, time.Duration(0), nil)
 
 	uh.Controller = cm
 	uh.Settings = &Settings{}
@@ -197,6 +199,51 @@ func TestStateProbeWithUpdateAvailableButAlreadyInstalled(t *testing.T) {
 	om.AssertExpectations(t)
 }
 
+func TestStateProbeTimeout(t *testing.T) {
+	aim := &activeinactivemock.ActiveInactiveMock{}
+	om := &objectmock.ObjectMock{}
+	cm := &controllermock.ControllerMock{}
+
+	mode := installmodes.RegisterInstallMode(installmodes.InstallMode{
+		Name:              "test",
+		CheckRequirements: func() error { return nil },
+		GetObject:         func() interface{} { return om },
+	})
+	defer mode.Unregister()
+
+	um, err := metadata.NewUpdateMetadata([]byte(validJSONMetadata))
+	assert.NoError(t, err)
+
+	apiClient := client.NewApiClient("address")
+	apiClient.CheckRedirect = nil
+
+	uh, err := newTestUpdateHub(NewProbeState(apiClient), aim)
+	assert.NoError(t, err)
+
+	uh.Controller = cm
+
+	uh.Store.Remove(uh.Settings.RuntimeSettingsPath)
+
+	sha256sum := sha256.Sum256([]byte(validJSONMetadata))
+	signature, _ := rsa.SignPKCS1v15(rand.Reader, testPrivateKey, crypto.SHA256, sha256sum[:])
+
+	timeoutReached := false
+	cm.On("ProbeUpdate", apiClient, 0).Run(func(args mock.Arguments) {
+		timeoutReached = true
+	}).Return((*metadata.UpdateMetadata)(nil), []byte{}, time.Duration(0), errs.Wrap(&testTimeoutError{}, "download update failed")).Once()
+	cm.On("ProbeUpdate", apiClient, 1).Return(um, signature, time.Duration(0), nil).Once()
+
+	next, _ := uh.GetState().Handle(uh)
+
+	assert.True(t, timeoutReached)
+
+	assert.Equal(t, NewDownloadingState(apiClient, um, &ProgressTrackerImpl{}), next)
+
+	aim.AssertExpectations(t)
+	om.AssertExpectations(t)
+	cm.AssertExpectations(t)
+}
+
 func TestStateProbeToMap(t *testing.T) {
 	state := NewProbeState(client.NewApiClient("address"))
 

updatehub/updatehub.go

@@ -278,12 +278,12 @@ func (uh *UpdateHub) ProcessCurrentState() State {
 }
 
 type Controller interface {
-	ProbeUpdate(*client.ApiClient, int) (*metadata.UpdateMetadata, []byte, time.Duration)
+	ProbeUpdate(*client.ApiClient, int) (*metadata.UpdateMetadata, []byte, time.Duration, error)
 	DownloadUpdate(*client.ApiClient, *metadata.UpdateMetadata, <-chan bool, chan<- int) error
 	InstallUpdate(*metadata.UpdateMetadata, chan<- int) error
 }
 
-func (uh *UpdateHub) ProbeUpdate(apiClient *client.ApiClient, retries int) (*metadata.UpdateMetadata, []byte, time.Duration) {
+func (uh *UpdateHub) ProbeUpdate(apiClient *client.ApiClient, retries int) (*metadata.UpdateMetadata, []byte, time.Duration, error) {
 	var data struct {
 		Retries int `json:"retries"`
 		metadata.FirmwareMetadata
@@ -299,18 +299,18 @@ func (uh *UpdateHub) ProbeUpdate(apiClient *client.ApiClient, retries int) (*met
 	updateMetadata, signature, extraPoll, err := uh.Updater.ProbeUpdate(apiClient.Request(), client.UpgradesEndpoint, data)
 	if err != nil {
 		uh.Store.Remove(updateMetadataPath)
-		return nil, nil, -1
+		return nil, nil, -1, err
 	}
 
 	if updateMetadata == nil || updateMetadata.(*metadata.UpdateMetadata) == nil {
 		uh.Store.Remove(updateMetadataPath)
-		return nil, signature, extraPoll
+		return nil, signature, extraPoll, nil
 	}
 
 	um := updateMetadata.(*metadata.UpdateMetadata)
 	afero.WriteFile(uh.Store, updateMetadataPath, um.RawBytes, 0644)
 
-	return um, signature, extraPoll
+	return um, signature, extraPoll, nil
 }
 
 // it is recommended to use a buffered channel for "progressChan" to ensure no progress event is lost

updatehub/updatehub_test.go

@@ -618,11 +618,12 @@ func TestUpdateHubProbeUpdate(t *testing.T) {
 			um.On("ProbeUpdate", apiClient.Request(), client.UpgradesEndpoint, data).Return(expectedUpdateMetadata, expectedSignature, tc.extraPoll, nil)
 			uh.Updater = um
 
-			updateMetadata, signature, extraPoll := uh.ProbeUpdate(apiClient, 0)
+			updateMetadata, signature, extraPoll, err := uh.ProbeUpdate(apiClient, 0)
 
 			assert.Equal(t, expectedUpdateMetadata, updateMetadata)
 			assert.Equal(t, expectedSignature, signature)
 			assert.Equal(t, tc.extraPoll, extraPoll)
+			assert.Nil(t, err)
 			um.AssertExpectations(t)
 
 			if tc.updateMetadata == "" {
@@ -676,11 +677,12 @@ func TestUpdateHubProbeUpdateWithNilUpdateMetadata(t *testing.T) {
 
 	uh.Updater = um
 
-	updateMetadata, signature, extraPoll := uh.ProbeUpdate(apiClient, 0)
+	updateMetadata, signature, extraPoll, err := uh.ProbeUpdate(apiClient, 0)
 
 	assert.Equal(t, (*metadata.UpdateMetadata)(nil), updateMetadata)
 	assert.Equal(t, []byte{}, signature)
 	assert.Equal(t, time.Duration(3000), extraPoll)
+	assert.Nil(t, err)
 	um.AssertExpectations(t)
 
 	fileExists, err := afero.Exists(uh.Store, updateMetadataPath)