fixe

skip-mev · Feb 7, 2025 · 9f9e21d · 9f9e21d
1 parent 99438c9
commit 9f9e21d
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 58 deletions.
diff --git a/core/provider/digitalocean/droplet.go b/core/provider/digitalocean/droplet.go
@@ -3,7 +3,6 @@ package digitalocean
 import (
 	"context"
 	"fmt"
-	"strings"
 	"time"
 
 	"github.com/pkg/errors"
@@ -58,34 +57,6 @@ func (p *Provider) CreateDroplet(ctx context.Context, definition provider.TaskDe
 	return droplet, nil
 }
 
-func (t *Task) waitForSSHClient(ctx context.Context) error {
-	start := time.Now()
-
-	err := util.WaitForCondition(ctx, time.Second*600, time.Millisecond*300, func() (bool, error) {
-		_, err := t.getDropletSSHClient(ctx)
-
-		if err != nil {
-			if strings.Contains(err.Error(), "connection refused") {
-				t.logger.Debug("connection refused", zap.String("task", t.GetState().Name))
-				return false, nil
-			}
-			return false, err
-		}
-
-		return true, nil
-	})
-
-	if err != nil {
-		return errors.Wrap(err, "failed to wait for ssh in droplet to become active")
-	}
-
-	end := time.Now()
-
-	t.logger.Info("droplet's ssh daemon is ready after", zap.String("name", t.GetState().Name), zap.Duration("startup_time", end.Sub(start)))
-
-	return nil
-}
-
 func (t *Task) waitForDockerStart(ctx context.Context) error {
 	start := time.Now()
 

diff --git a/core/provider/digitalocean/provider.go b/core/provider/digitalocean/provider.go
@@ -163,10 +163,6 @@ func (p *Provider) CreateTask(ctx context.Context, definition provider.TaskDefin
 		return nil, fmt.Errorf("failed to wait for tailscale peer: %w", err)
 	}
 
-	if err := task.waitForSSHClient(ctx); err != nil {
-		return nil, fmt.Errorf("failed to wait for SSH client: %w", err)
-	}
-
 	ip, err := task.GetIP(ctx)
 
 	if err != nil {

diff --git a/core/provider/digitalocean/provider_test.go b/core/provider/digitalocean/provider_test.go
@@ -52,8 +52,6 @@ func setupTestProvider(t *testing.T, ctx context.Context) (*Provider, *mocks.Moc
 	mockTailscaleServer := clientmocks.NewMockTailscaleServer(t)
 	mockTailscaleClient := clientmocks.NewMockTailscaleLocalClient(t)
 
-	mockTailscaleClient.On("Status", ctx).Return(generateTailscaleStatus(t, "test-task", "1.2.3.4"), nil)
-
 	mockTailscale := TailscaleSettings{
 		Server:      mockTailscaleServer,
 		LocalClient: mockTailscaleClient,
@@ -87,8 +85,6 @@ func setupTestProvider(t *testing.T, ctx context.Context) (*Provider, *mocks.Moc
 
 	mockDO.On("CreateTag", ctx, mock.Anything).Return(&godo.Tag{Name: "test-tag"}, nil)
 	mockDO.On("CreateFirewall", ctx, mock.Anything).Return(&godo.Firewall{ID: "test-firewall"}, nil)
-	mockDO.On("GetKeyByFingerprint", ctx, mock.AnythingOfType("string")).Return(nil, nil)
-	mockDO.On("CreateKey", ctx, mock.Anything).Return(&godo.Key{}, nil)
 
 	mockDockerClients := map[string]clients.DockerClient{
 		"test-task": mockDocker,
@@ -97,6 +93,8 @@ func setupTestProvider(t *testing.T, ctx context.Context) (*Provider, *mocks.Moc
 	p, err := NewProviderWithClient(ctx, "test-provider", mockDO, mockTailscale, WithDockerClients(mockDockerClients), WithLogger(logger))
 	require.NoError(t, err)
 
+	mockTailscaleClient.On("Status", ctx).Return(generateTailscaleStatus(t, fmt.Sprintf("%s-test-task", p.GetState().PetriTag), "1.2.3.4"), nil)
+
 	droplet := &godo.Droplet{
 		ID: 123,
 		Networks: &godo.Networks{
@@ -181,8 +179,6 @@ func setupValidationTestProvider(t *testing.T, ctx context.Context) *Provider {
 
 	mockDO.On("CreateTag", ctx, mock.Anything).Return(&godo.Tag{Name: "test-tag"}, nil)
 	mockDO.On("CreateFirewall", ctx, mock.Anything).Return(&godo.Firewall{ID: "test-firewall"}, nil)
-	mockDO.On("GetKeyByFingerprint", ctx, mock.AnythingOfType("string")).Return(nil, nil)
-	mockDO.On("CreateKey", ctx, mock.Anything).Return(&godo.Key{}, nil)
 
 	p, err := NewProviderWithClient(ctx, "test-provider", mockDO, mockTailscale, WithLogger(logger))
 	require.NoError(t, err)
@@ -332,23 +328,22 @@ func TestConcurrentTaskCreationAndCleanup(t *testing.T) {
 		}, nil).Maybe()
 		mockDocker.On("Close").Return(nil).Once()
 
-		mockStatus := ipnstate.PeerStatus{HostName: fmt.Sprintf("test-task-%d", i), TailscaleIPs: []netip.Addr{netip.MustParseAddr(fmt.Sprintf("1.2.3.%d", i+1))}}
-		mockStatuses[key.NewNode().Public()] = &mockStatus
 	}
 
-	mockTailscaleClient.On("Status", ctx).Return(&ipnstate.Status{Peer: mockStatuses}, nil)
-
 	mockDO.On("CreateTag", ctx, mock.Anything).Return(&godo.Tag{Name: "test-tag"}, nil)
 
 	mockDO.On("CreateFirewall", ctx, mock.Anything).Return(&godo.Firewall{ID: "test-firewall"}, nil)
 
-	mockDO.On("GetKeyByFingerprint", ctx, mock.AnythingOfType("string")).
-		Return(nil, nil)
-	mockDO.On("CreateKey", ctx, mock.Anything).Return(&godo.Key{}, nil)
-
 	p, err := NewProviderWithClient(ctx, "test-provider", mockDO, mockTailscale, WithDockerClients(mockDockerClients), WithLogger(logger))
 	require.NoError(t, err)
 
+	for i := 0; i < 10; i++ {
+		mockStatus := ipnstate.PeerStatus{HostName: fmt.Sprintf("%s-test-task-%d", p.GetState().PetriTag, i), TailscaleIPs: []netip.Addr{netip.MustParseAddr(fmt.Sprintf("1.2.3.%d", i+1))}}
+		mockStatuses[key.NewNode().Public()] = &mockStatus
+	}
+
+	mockTailscaleClient.On("Status", ctx).Return(&ipnstate.Status{Peer: mockStatuses}, nil)
+
 	numTasks := 10
 	var wg sync.WaitGroup
 	errors := make(chan error, numTasks)
@@ -380,7 +375,6 @@ func TestConcurrentTaskCreationAndCleanup(t *testing.T) {
 
 	mockDO.On("DeleteDropletByTag", ctx, mock.AnythingOfType("string")).Return(nil).Once()
 	mockDO.On("DeleteFirewall", ctx, mock.AnythingOfType("string")).Return(nil).Once()
-	mockDO.On("DeleteKeyByFingerprint", ctx, mock.AnythingOfType("string")).Return(nil).Once()
 	mockDO.On("DeleteTag", ctx, mock.AnythingOfType("string")).Return(nil).Once()
 
 	for i := 0; i < numTasks; i++ {
@@ -504,8 +498,6 @@ func TestProviderSerialization(t *testing.T) {
 	mockTailscaleServer := clientmocks.NewMockTailscaleServer(t)
 	mockTailscaleClient := clientmocks.NewMockTailscaleLocalClient(t)
 
-	mockTailscaleClient.On("Status", ctx).Return(generateTailscaleStatus(t, "test-task", "1.2.3.4"), nil)
-
 	mockTailscale := TailscaleSettings{
 		Server:      mockTailscaleServer,
 		LocalClient: mockTailscaleClient,
@@ -515,8 +507,6 @@ func TestProviderSerialization(t *testing.T) {
 
 	mockDO.On("CreateTag", ctx, mock.Anything).Return(&godo.Tag{Name: "petri-droplet-test"}, nil)
 	mockDO.On("CreateFirewall", ctx, mock.Anything).Return(&godo.Firewall{ID: "test-firewall"}, nil)
-	mockDO.On("GetKeyByFingerprint", ctx, mock.AnythingOfType("string")).Return(nil, nil)
-	mockDO.On("CreateKey", ctx, mock.Anything).Return(&godo.Key{}, nil)
 
 	mockDockerClients := map[string]clients.DockerClient{
 		"test-task": mockDocker,
@@ -525,6 +515,8 @@ func TestProviderSerialization(t *testing.T) {
 	p1, err := NewProviderWithClient(ctx, "test-provider", mockDO, mockTailscale, WithDockerClients(mockDockerClients), WithLogger(zap.NewExample()))
 	require.NoError(t, err)
 
+	mockTailscaleClient.On("Status", ctx).Return(generateTailscaleStatus(t, fmt.Sprintf("%s-test-task", p1.GetState().PetriTag), "1.2.3.4"), nil)
+
 	droplet := &godo.Droplet{
 		ID: 123,
 		Networks: &godo.Networks{

diff --git a/core/provider/digitalocean/task_test.go b/core/provider/digitalocean/task_test.go
@@ -395,9 +395,10 @@ func TestTaskGetIP(t *testing.T) {
 
 	task := &Task{
 		state: &TaskState{
-			ID:           strconv.Itoa(droplet.ID),
-			Name:         "test-task",
-			ProviderName: "test-provider",
+			TailscaleHostname: "test-task",
+			ID:                strconv.Itoa(droplet.ID),
+			Name:              "test-task",
+			ProviderName:      "test-provider",
 		},
 		logger:            logger,
 		dockerClient:      mockDocker,
@@ -759,9 +760,10 @@ func TestTaskExposingPort(t *testing.T) {
 
 	task := &Task{
 		state: &TaskState{
-			ID:           strconv.Itoa(droplet.ID),
-			Name:         "test-task",
-			ProviderName: "test-provider",
+			TailscaleHostname: "test-task",
+			ID:                strconv.Itoa(droplet.ID),
+			Name:              "test-task",
+			ProviderName:      "test-provider",
 			Definition: provider.TaskDefinition{
 				Name: "test-task",
 				Image: provider.ImageDefinition{