Skip to content

Commit 22eb487

Browse files
ChaitanyaKulkarni28ChaitanyaKulkarni28
authored
Update metadata script runner, add tests (#357)
1 parent a13f564 commit 22eb487

File tree

2 files changed

+187
-37
lines changed

2 files changed

+187
-37
lines changed

‎google_metadata_script_runner/main.go

Lines changed: 32 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@
1717
package main
1818

1919
// TODO: compare log outputs in this utility to linux.
20-
// TODO: standardize and consolidate retries.
2120

2221
import (
2322
"bufio"
@@ -40,6 +39,7 @@ import (
4039
"cloud.google.com/go/storage"
4140
"github.com/GoogleCloudPlatform/guest-agent/google_guest_agent/cfg"
4241
"github.com/GoogleCloudPlatform/guest-agent/metadata"
42+
"github.com/GoogleCloudPlatform/guest-agent/retry"
4343
"github.com/GoogleCloudPlatform/guest-agent/utils"
4444
"github.com/GoogleCloudPlatform/guest-logging-go/logger"
4545
)
@@ -79,10 +79,13 @@ var (
7979
// https://commondatastorage.googleapis.com/<bucket>/<object>
8080
gsHTTPRegex3 = regexp.MustCompile(fmt.Sprintf(`^http[s]?://(?:commondata)?storage\.googleapis\.com/%s/%s$`, bucket, object))
8181

82+
// testStorageClient is used to override GCS client in unit tests.
8283
testStorageClient *storage.Client
8384

8485
client metadata.MDSClientInterface
8586
version string
87+
// defaultRetryPolicy is default policy to retry up to 3 times, only wait 1 second between retries.
88+
defaultRetryPolicy = retry.Policy{MaxAttempts: 3, BackoffFactor: 1, Jitter: time.Second}
8689
)
8790

8891
func init() {
@@ -103,36 +106,35 @@ func downloadGSURL(ctx context.Context, bucket, object string, file *os.File) er
103106
}
104107
defer client.Close()
105108

106-
r, err := client.Bucket(bucket).Object(object).NewReader(ctx)
109+
r, err := retry.RunWithResponse(ctx, defaultRetryPolicy, func() (*storage.Reader, error) {
110+
r, err := client.Bucket(bucket).Object(object).NewReader(ctx)
111+
return r, err
112+
})
107113
if err != nil {
108-
return fmt.Errorf("error reading object %q: %v", object, err)
114+
return err
109115
}
110116
defer r.Close()
111117

112118
_, err = io.Copy(file, r)
113119
return err
114120
}
115121

116-
func downloadURL(url string, file *os.File) error {
117-
// Retry up to 3 times, only wait 1 second between retries.
118-
var res *http.Response
119-
var err error
120-
for i := 1; ; i++ {
121-
res, err = http.Get(url)
122-
if err != nil && i > 3 {
123-
return err
122+
func downloadURL(ctx context.Context, url string, file *os.File) error {
123+
res, err := retry.RunWithResponse(context.Background(), defaultRetryPolicy, func() (*http.Response, error) {
124+
res, err := http.Get(url)
125+
if err != nil {
126+
return res, err
124127
}
125-
if err == nil {
126-
break
128+
if res.StatusCode != http.StatusOK {
129+
return nil, fmt.Errorf("GET %q, bad status: %s", url, res.Status)
127130
}
128-
time.Sleep(1 * time.Second)
131+
return res, nil
132+
})
133+
if err != nil {
134+
return err
129135
}
130136
defer res.Body.Close()
131137

132-
if res.StatusCode != http.StatusOK {
133-
return fmt.Errorf("GET %q, bad status: %s", url, res.Status)
134-
}
135-
136138
_, err = io.Copy(file, res.Body)
137139
return err
138140
}
@@ -142,34 +144,27 @@ func downloadScript(ctx context.Context, path string, file *os.File) error {
142144
// particularly once a system is promoted to a domain controller.
143145
// Try to lookup storage.googleapis.com and sleep for up to 100s if
144146
// we get an error.
145-
// TODO: do we need to do this on every script?
146-
for i := 0; i < 20; i++ {
147-
if _, err := net.LookupHost(storageURL); err == nil {
148-
break
149-
}
150-
time.Sleep(5 * time.Second)
147+
policy := retry.Policy{MaxAttempts: 20, BackoffFactor: 1, Jitter: time.Second * 5}
148+
err := retry.Run(ctx, policy, func() error {
149+
_, err := net.LookupHost(storageURL)
150+
return err
151+
})
152+
if err != nil {
153+
return fmt.Errorf("%q lookup failed, err: %+v", storageURL, err)
151154
}
155+
152156
bucket, object := parseGCS(path)
153157
if bucket != "" && object != "" {
154-
// TODO: why is this retry outer, but downloadURL retry is inner?
155-
// Retry up to 3 times, only wait 1 second between retries.
156-
for i := 1; ; i++ {
157-
err := downloadGSURL(ctx, bucket, object, file)
158-
if err == nil {
159-
return nil
160-
}
161-
if err != nil && i > 3 {
162-
logger.Infof("Failed to download GCS path: %v", err)
163-
break
164-
}
165-
time.Sleep(1 * time.Second)
158+
err = downloadGSURL(ctx, bucket, object, file)
159+
if err != nil {
160+
logger.Infof("Failed to download object [%s] from GCS bucket [%s], err: %+v", object, bucket, err)
166161
}
167162
logger.Infof("Trying unauthenticated download")
168163
path = fmt.Sprintf("https://%s/%s/%s", storageURL, bucket, object)
169164
}
170165

171166
// Fall back to an HTTP GET of the URL.
172-
return downloadURL(path, file)
167+
return downloadURL(ctx, path, file)
173168
}
174169

175170
func parseGCS(path string) (string, string) {

‎google_metadata_script_runner/main_test.go

Lines changed: 155 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,13 +17,20 @@ package main
1717
import (
1818
"context"
1919
"fmt"
20+
"net/http"
21+
"net/http/httptest"
2022
"net/url"
2123
"os"
24+
"path/filepath"
2225
"reflect"
26+
"strings"
2327
"testing"
28+
"time"
2429

30+
"cloud.google.com/go/storage"
2531
"github.com/GoogleCloudPlatform/guest-agent/google_guest_agent/cfg"
2632
"github.com/GoogleCloudPlatform/guest-agent/metadata"
33+
"google.golang.org/api/option"
2734
)
2835

2936
func TestMain(m *testing.M) {
@@ -297,3 +304,151 @@ func TestGetWantedKeysError(t *testing.T) {
297304
})
298305
}
299306
}
307+
308+
func TestDownloadURL(t *testing.T) {
309+
ctx := context.Background()
310+
ctr := make(map[string]int)
311+
// No need to wait longer, override for testing.
312+
defaultRetryPolicy.Jitter = time.Millisecond
313+
314+
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
315+
// /retry should succeed within 2 retries; /fail should always fail.
316+
if (r.URL.Path == "/retry" && ctr["/retry"] != 1) || strings.Contains(r.URL.Path, "fail") {
317+
w.WriteHeader(400)
318+
}
319+
320+
fmt.Fprintf(w, r.URL.Path)
321+
ctr[r.URL.Path] = ctr[r.URL.Path] + 1
322+
}))
323+
defer server.Close()
324+
325+
tests := []struct {
326+
name string
327+
key string
328+
wantErr bool
329+
retries int
330+
}{
331+
{
332+
name: "succeed_immediately",
333+
key: "/immediate_download",
334+
wantErr: false,
335+
retries: 1,
336+
},
337+
{
338+
name: "succeed_after_retry",
339+
key: "/retry",
340+
wantErr: false,
341+
retries: 2,
342+
},
343+
{
344+
name: "fail_retry_exhaust",
345+
key: "/fail",
346+
wantErr: true,
347+
retries: 3,
348+
},
349+
}
350+
for _, tt := range tests {
351+
t.Run(tt.name, func(t *testing.T) {
352+
f, err := os.OpenFile(filepath.Join(t.TempDir(), tt.name), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0755)
353+
if err != nil {
354+
t.Fatalf("Failed to setup test file: %v", err)
355+
}
356+
defer f.Close()
357+
url := server.URL + tt.key
358+
if err := downloadURL(ctx, url, f); (err != nil) != tt.wantErr {
359+
t.Errorf("downloadURL(ctx, %s, %s) error = [%v], wantErr %t", url, f.Name(), err, tt.wantErr)
360+
}
361+
362+
if !tt.wantErr {
363+
gotBytes, err := os.ReadFile(f.Name())
364+
if err != nil {
365+
t.Errorf("failed to read output file %q, with error: %v", f.Name(), err)
366+
}
367+
if string(gotBytes) != tt.key {
368+
t.Errorf("downloadURL(ctx, %s, %s) wrote = [%s], want [%s]", url, f.Name(), string(gotBytes), tt.key)
369+
}
370+
}
371+
372+
if ctr[tt.key] != tt.retries {
373+
t.Errorf("downloadURL(ctx, %s, %s) retried [%d] times, should have returned after [%d] retries", url, f.Name(), ctr[tt.key], tt.retries)
374+
}
375+
})
376+
}
377+
}
378+
379+
func TestDownloadGSURL(t *testing.T) {
380+
ctx := context.Background()
381+
ctr := make(map[string]int)
382+
// No need to wait longer, override for testing.
383+
defaultRetryPolicy.Jitter = time.Millisecond
384+
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
385+
// Fake error for invalid object request.
386+
if strings.Contains(r.URL.Path, "invalid") {
387+
w.WriteHeader(404)
388+
}
389+
fmt.Fprintf(w, r.URL.Path)
390+
ctr[r.URL.Path] = ctr[r.URL.Path] + 1
391+
}))
392+
defer server.Close()
393+
394+
var err error
395+
httpClient := &http.Client{Transport: &http.Transport{}}
396+
testStorageClient, err = storage.NewClient(ctx, option.WithHTTPClient(httpClient), option.WithEndpoint(server.URL))
397+
if err != nil {
398+
t.Fatalf("Failed to setup test storage client, err: %+v", err)
399+
}
400+
defer testStorageClient.Close()
401+
402+
tests := []struct {
403+
name string
404+
bucket string
405+
object string
406+
wantErr bool
407+
retries int
408+
}{
409+
{
410+
name: "valid_object",
411+
bucket: "valid",
412+
object: "obj1",
413+
wantErr: false,
414+
retries: 1,
415+
},
416+
{
417+
name: "invalid_object",
418+
bucket: "invalid",
419+
object: "obj1",
420+
wantErr: true,
421+
retries: 3,
422+
},
423+
}
424+
for _, tt := range tests {
425+
t.Run(tt.name, func(t *testing.T) {
426+
f, err := os.OpenFile(filepath.Join(t.TempDir(), tt.name), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0755)
427+
if err != nil {
428+
t.Fatalf("Failed to setup test file: %v", err)
429+
}
430+
defer f.Close()
431+
432+
if err := downloadGSURL(ctx, tt.bucket, tt.object, f); (err != nil) != tt.wantErr {
433+
t.Errorf("downloadGSURL(ctx, %s, %s, %s) error = [%+v], wantErr %t", tt.bucket, tt.object, f.Name(), err, tt.wantErr)
434+
}
435+
436+
want := fmt.Sprintf("/%s/%s", tt.bucket, tt.object)
437+
438+
if !tt.wantErr {
439+
gotBytes, err := os.ReadFile(f.Name())
440+
if err != nil {
441+
t.Errorf("failed to read output file %q, with error: %v", f.Name(), err)
442+
}
443+
444+
if string(gotBytes) != want {
445+
t.Errorf("downloadGSURL(ctx, %s, %s, %s) wrote = [%s], want [%s]", tt.bucket, tt.object, f.Name(), string(gotBytes), want)
446+
}
447+
}
448+
449+
if ctr[want] != tt.retries {
450+
t.Errorf("downloadGSURL(ctx, %s, %s, %s) retried [%d] times, should have returned after [%d] retries", tt.bucket, tt.object, f.Name(), ctr[want], tt.retries)
451+
}
452+
})
453+
}
454+
}

0 commit comments

Comments
 (0)