fix: add timeout to get workspace info from provider (#1058)

Signed-off-by: Ivan Dagelic <[email protected]>

internal/testing/server/workspaces/mocks/provisioner.go

@@ -6,6 +6,8 @@
 package mocks
 
 import (
+	"context"
+
 	"github.com/daytonaio/daytona/pkg/containerregistry"
 	"github.com/daytonaio/daytona/pkg/gitprovider"
 	"github.com/daytonaio/daytona/pkg/provider"
@@ -42,8 +44,8 @@ func (p *mockProvisioner) DestroyWorkspace(workspace *workspace.Workspace, targe
 	return args.Error(0)
 }
 
-func (p *mockProvisioner) GetWorkspaceInfo(w *workspace.Workspace, target *provider.ProviderTarget) (*workspace.WorkspaceInfo, error) {
-	args := p.Called(w, target)
+func (p *mockProvisioner) GetWorkspaceInfo(ctx context.Context, w *workspace.Workspace, target *provider.ProviderTarget) (*workspace.WorkspaceInfo, error) {
+	args := p.Called(ctx, w, target)
 	return args.Get(0).(*workspace.WorkspaceInfo), args.Error(1)
 }
 

pkg/api/controllers/workspace/workspace.go

@@ -64,7 +64,7 @@ func ListWorkspaces(ctx *gin.Context) {
 
 	server := server.GetInstance(nil)
 
-	workspaceList, err := server.WorkspaceService.ListWorkspaces(verbose)
+	workspaceList, err := server.WorkspaceService.ListWorkspaces(ctx.Request.Context(), verbose)
 	if err != nil {
 		ctx.AbortWithError(http.StatusInternalServerError, fmt.Errorf("failed to list workspaces: %w", err))
 		return

pkg/provisioner/info.go

@@ -4,18 +4,40 @@
 package provisioner
 
 import (
+	"context"
+
 	"github.com/daytonaio/daytona/pkg/provider"
 	"github.com/daytonaio/daytona/pkg/workspace"
 )
 
-func (p *Provisioner) GetWorkspaceInfo(workspace *workspace.Workspace, target *provider.ProviderTarget) (*workspace.WorkspaceInfo, error) {
-	targetProvider, err := p.providerManager.GetProvider(target.ProviderInfo.Name)
-	if err != nil {
-		return nil, err
-	}
+type InfoResult struct {
+	Info *workspace.WorkspaceInfo
+	Err  error
+}
+
+// Gets the workspace info from the provider - the context is used to cancel the request if it takes too long
+func (p *Provisioner) GetWorkspaceInfo(ctx context.Context, ws *workspace.Workspace, target *provider.ProviderTarget) (*workspace.WorkspaceInfo, error) {
+	ch := make(chan InfoResult, 1)
 
-	return (*targetProvider).GetWorkspaceInfo(&provider.WorkspaceRequest{
-		TargetOptions: target.Options,
-		Workspace:     workspace,
-	})
+	go func() {
+		targetProvider, err := p.providerManager.GetProvider(target.ProviderInfo.Name)
+		if err != nil {
+			ch <- InfoResult{nil, err}
+			return
+		}
+
+		info, err := (*targetProvider).GetWorkspaceInfo(&provider.WorkspaceRequest{
+			TargetOptions: target.Options,
+			Workspace:     ws,
+		})
+
+		ch <- InfoResult{info, err}
+	}()
+
+	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	case data := <-ch:
+		return data.Info, data.Err
+	}
 }

pkg/provisioner/provisioner.go

@@ -4,6 +4,8 @@
 package provisioner
 
 import (
+	"context"
+
 	"github.com/daytonaio/daytona/pkg/containerregistry"
 	"github.com/daytonaio/daytona/pkg/gitprovider"
 	"github.com/daytonaio/daytona/pkg/provider"
@@ -17,7 +19,7 @@ type IProvisioner interface {
 	CreateWorkspace(workspace *workspace.Workspace, target *provider.ProviderTarget) error
 	DestroyProject(project *project.Project, target *provider.ProviderTarget) error
 	DestroyWorkspace(workspace *workspace.Workspace, target *provider.ProviderTarget) error
-	GetWorkspaceInfo(workspace *workspace.Workspace, target *provider.ProviderTarget) (*workspace.WorkspaceInfo, error)
+	GetWorkspaceInfo(ctx context.Context, workspace *workspace.Workspace, target *provider.ProviderTarget) (*workspace.WorkspaceInfo, error)
 	StartProject(project *project.Project, target *provider.ProviderTarget) error
 	StartWorkspace(workspace *workspace.Workspace, target *provider.ProviderTarget) error
 	StopProject(project *project.Project, target *provider.ProviderTarget) error

pkg/server/purge.go

@@ -38,7 +38,7 @@ func (s *Server) Purge(ctx context.Context, force bool) []error {
 		return []error{err}
 	}
 
-	workspaces, err := s.WorkspaceService.ListWorkspaces(false)
+	workspaces, err := s.WorkspaceService.ListWorkspaces(ctx, false)
 	if err != nil {
 		s.trackPurgeError(ctx, force, err)
 		if !force {

pkg/server/workspaces/get.go

@@ -5,29 +5,54 @@ package workspaces
 
 import (
 	"context"
+	"errors"
+	"fmt"
+	"time"
 
+	"github.com/daytonaio/daytona/pkg/provisioner"
 	"github.com/daytonaio/daytona/pkg/server/workspaces/dto"
+	log "github.com/sirupsen/logrus"
 )
 
 func (s *WorkspaceService) GetWorkspace(ctx context.Context, workspaceId string) (*dto.WorkspaceDTO, error) {
-	workspace, err := s.workspaceStore.Find(workspaceId)
+	ws, err := s.workspaceStore.Find(workspaceId)
 	if err != nil {
 		return nil, ErrWorkspaceNotFound
 	}
 
-	target, err := s.targetStore.Find(workspace.Target)
-	if err != nil {
-		return nil, err
+	response := dto.WorkspaceDTO{
+		Workspace: *ws,
 	}
 
-	workspaceInfo, err := s.provisioner.GetWorkspaceInfo(workspace, target)
+	target, err := s.targetStore.Find(ws.Target)
 	if err != nil {
 		return nil, err
 	}
 
-	response := dto.WorkspaceDTO{
-		Workspace: *workspace,
-		Info:      workspaceInfo,
+	ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
+	defer cancel()
+
+	resultCh := make(chan provisioner.InfoResult, 1)
+
+	go func() {
+		workspaceInfo, err := s.provisioner.GetWorkspaceInfo(ctx, ws, target)
+		resultCh <- provisioner.InfoResult{Info: workspaceInfo, Err: err}
+	}()
+
+	select {
+	case res := <-resultCh:
+		if res.Err != nil {
+			log.Error(fmt.Errorf("failed to get workspace info for %s: %v", ws.Name, res.Err))
+			return nil, res.Err
+		}
+
+		response.Info = res.Info
+	case <-ctx.Done():
+		if errors.Is(ctx.Err(), context.DeadlineExceeded) {
+			log.Warn(fmt.Sprintf("timeout getting workspace info for %s", ws.Name))
+		} else {
+			log.Warn(fmt.Sprintf("cancelled getting workspace info for %s", ws.Name))
+		}
 	}
 
 	return &response, nil

pkg/server/workspaces/list.go

@@ -4,41 +4,76 @@
 package workspaces
 
 import (
+	"context"
+	"errors"
 	"fmt"
+	"sync"
+	"time"
 
+	"github.com/daytonaio/daytona/pkg/provisioner"
 	"github.com/daytonaio/daytona/pkg/server/workspaces/dto"
 	"github.com/daytonaio/daytona/pkg/workspace"
 	log "github.com/sirupsen/logrus"
 )
 
-func (s *WorkspaceService) ListWorkspaces(verbose bool) ([]dto.WorkspaceDTO, error) {
+func (s *WorkspaceService) ListWorkspaces(ctx context.Context, verbose bool) ([]dto.WorkspaceDTO, error) {
 	workspaces, err := s.workspaceStore.List()
 	if err != nil {
 		return nil, err
 	}
 
+	var wg sync.WaitGroup
 	response := []dto.WorkspaceDTO{}
 
-	for _, w := range workspaces {
-		var workspaceInfo *workspace.WorkspaceInfo
-		if verbose {
+	for i, w := range workspaces {
+		if !verbose {
+			response = append(response, dto.WorkspaceDTO{Workspace: *w})
+			continue
+		}
+
+		wg.Add(1)
+		go func(i int, w *workspace.Workspace) {
+			defer wg.Done()
+
+			workspaceDto := dto.WorkspaceDTO{Workspace: *w}
+
 			target, err := s.targetStore.Find(w.Target)
 			if err != nil {
 				log.Error(fmt.Errorf("failed to get target for %s", w.Target))
-				continue
+				return
 			}
 
-			workspaceInfo, err = s.provisioner.GetWorkspaceInfo(w, target)
-			if err != nil {
-				log.Error(fmt.Errorf("failed to get workspace info for %s", w.Name))
-			}
-		}
+			ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
+			defer cancel()
+
+			resultCh := make(chan provisioner.InfoResult, 1)
+
+			go func() {
+				workspaceInfo, err := s.provisioner.GetWorkspaceInfo(ctx, w, target)
+				resultCh <- provisioner.InfoResult{Info: workspaceInfo, Err: err}
+			}()
 
-		response = append(response, dto.WorkspaceDTO{
-			Workspace: *w,
-			Info:      workspaceInfo,
-		})
+			select {
+			case res := <-resultCh:
+				if res.Err != nil {
+					log.Error(fmt.Errorf("failed to get workspace info for %s: %v", w.Name, res.Err))
+					response = append(response, workspaceDto)
+					return
+				}
+
+				workspaceDto.Info = res.Info
+				response = append(response, workspaceDto)
+			case <-ctx.Done():
+				if errors.Is(ctx.Err(), context.DeadlineExceeded) {
+					log.Warn(fmt.Sprintf("timeout getting workspace info for %s", w.Name))
+				} else {
+					log.Warn(fmt.Sprintf("cancelled getting workspace info for %s", w.Name))
+				}
+				response = append(response, workspaceDto)
+			}
+		}(i, w)
 	}
 
+	wg.Wait()
 	return response, nil
 }

pkg/server/workspaces/service.go

@@ -27,7 +27,7 @@ type IWorkspaceService interface {
 	GetWorkspace(ctx context.Context, workspaceId string) (*dto.WorkspaceDTO, error)
 	GetWorkspaceLogReader(workspaceId string) (io.Reader, error)
 	GetProjectLogReader(workspaceId, projectName string) (io.Reader, error)
-	ListWorkspaces(verbose bool) ([]dto.WorkspaceDTO, error)
+	ListWorkspaces(ctx context.Context, verbose bool) ([]dto.WorkspaceDTO, error)
 	RemoveWorkspace(ctx context.Context, workspaceId string) error
 	ForceRemoveWorkspace(ctx context.Context, workspaceId string) error
 	SetProjectState(workspaceId string, projectName string, state *project.ProjectState) (*workspace.Workspace, error)

pkg/server/workspaces/service_test.go

@@ -158,7 +158,7 @@ func TestWorkspaceService(t *testing.T) {
 	})
 
 	t.Run("GetWorkspace", func(t *testing.T) {
-		provisioner.On("GetWorkspaceInfo", mock.Anything, &target).Return(&workspaceInfo, nil)
+		provisioner.On("GetWorkspaceInfo", mock.Anything, mock.Anything, &target).Return(&workspaceInfo, nil)
 
 		workspace, err := service.GetWorkspace(context.TODO(), createWorkspaceDto.Id)
 
@@ -176,9 +176,9 @@ func TestWorkspaceService(t *testing.T) {
 
 	t.Run("ListWorkspaces", func(t *testing.T) {
 		verbose := false
-		provisioner.On("GetWorkspaceInfo", mock.Anything, &target).Return(&workspaceInfo, nil)
+		provisioner.On("GetWorkspaceInfo", mock.Anything, mock.Anything, &target).Return(&workspaceInfo, nil)
 
-		workspaces, err := service.ListWorkspaces(verbose)
+		workspaces, err := service.ListWorkspaces(context.TODO(), verbose)
 
 		require.Nil(t, err)
 		require.Len(t, workspaces, 1)
@@ -190,9 +190,9 @@ func TestWorkspaceService(t *testing.T) {
 
 	t.Run("ListWorkspaces - verbose", func(t *testing.T) {
 		verbose := true
-		provisioner.On("GetWorkspaceInfo", mock.Anything, &target).Return(&workspaceInfo, nil)
+		provisioner.On("GetWorkspaceInfo", mock.Anything, mock.Anything, &target).Return(&workspaceInfo, nil)
 
-		workspaces, err := service.ListWorkspaces(verbose)
+		workspaces, err := service.ListWorkspaces(context.TODO(), verbose)
 
 		require.Nil(t, err)
 		require.Len(t, workspaces, 1)