Merge pull request #120 from skip-mev/nv/caching-droplet-ips

feat: allow caching of static droplet ips
skip-mev · Nov 5, 2024 · 8dd3c81 · 8dd3c81
2 parents b93f7f9 + 96a82ba
commit 8dd3c81
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 5 deletions.
diff --git a/core/provider/digitalocean/droplet.go b/core/provider/digitalocean/droplet.go
@@ -54,7 +54,7 @@ func (p *Provider) CreateDroplet(ctx context.Context, definition provider.TaskDe
 
 	start := time.Now()
 
-	err = util.WaitForCondition(ctx, time.Second*100, time.Millisecond*100, func() (bool, error) {
+	err = util.WaitForCondition(ctx, time.Second*300, time.Millisecond*300, func() (bool, error) {
 		d, _, err := p.doClient.Droplets.Get(ctx, droplet.ID)
 
 		if err != nil {
@@ -119,13 +119,16 @@ func (p *Provider) deleteDroplet(ctx context.Context, name string) error {
 	return nil
 }
 
-func (p *Provider) getDroplet(ctx context.Context, name string) (*godo.Droplet, error) {
+func (p *Provider) getDroplet(ctx context.Context, name string, returnOnCacheHit bool) (*godo.Droplet, error) {
 	cachedDroplet, ok := p.droplets.Load(name)
-
 	if !ok {
 		return nil, fmt.Errorf("could not find droplet %s", name)
 	}
 
+	if ok && returnOnCacheHit {
+		return cachedDroplet, nil
+	}
+
 	droplet, res, err := p.doClient.Droplets.Get(ctx, cachedDroplet.ID)
 
 	if err != nil {

diff --git a/core/provider/digitalocean/task.go b/core/provider/digitalocean/task.go
@@ -172,7 +172,7 @@ func (p *Provider) DestroyTask(ctx context.Context, taskName string) error {
 }
 
 func (p *Provider) GetTaskStatus(ctx context.Context, taskName string) (provider.TaskStatus, error) {
-	droplet, err := p.getDroplet(ctx, taskName)
+	droplet, err := p.getDroplet(ctx, taskName, false)
 
 	if err != nil {
 		return provider.TASK_STATUS_UNDEFINED, err
@@ -296,7 +296,7 @@ func (p *Provider) DownloadDir(ctx context.Context, s string, s2 string, s3 stri
 }
 
 func (p *Provider) GetIP(ctx context.Context, taskName string) (string, error) {
-	droplet, err := p.getDroplet(ctx, taskName)
+	droplet, err := p.getDroplet(ctx, taskName, true)
 
 	if err != nil {
 		return "", err