Make DNS resolver configurable

cmd/config.go

@@ -35,6 +35,7 @@ import (
 
 	"github.com/loadimpact/k6/lib"
 	"github.com/loadimpact/k6/lib/executor"
+	"github.com/loadimpact/k6/lib/types"
 	"github.com/loadimpact/k6/stats"
 	"github.com/loadimpact/k6/stats/cloud"
 	"github.com/loadimpact/k6/stats/csv"
@@ -259,6 +260,17 @@ func applyDefault(conf Config) Config {
 	if conf.Options.SummaryTrendStats == nil {
 		conf.Options.SummaryTrendStats = lib.DefaultSummaryTrendStats
 	}
+	defDNS := types.DefaultDNSConfig()
+	if !conf.DNS.TTL.Valid {
+		conf.DNS.TTL = defDNS.TTL
+	}
+	if !conf.DNS.Select.Valid {
+		conf.DNS.Select = defDNS.Select
+	}
+	if !conf.DNS.Policy.Valid {
+		conf.DNS.Policy = defDNS.Policy
+	}
+
 	return conf
 }
 

cmd/config_consolidation_test.go

@@ -139,7 +139,7 @@ type file struct {
 func getFS(files []file) afero.Fs {
 	fs := afero.NewMemMapFs()
 	for _, f := range files {
-		must(afero.WriteFile(fs, f.filepath, []byte(f.contents), 0644)) // modes don't matter in the afero.MemMapFs
+		must(afero.WriteFile(fs, f.filepath, []byte(f.contents), 0o644)) // modes don't matter in the afero.MemMapFs
 	}
 	return fs
 }
@@ -214,11 +214,13 @@ func getConfigConsolidationTestCases() []configConsolidationTestCase {
 		{opts{cli: []string{"-u", "3", "-d", "30s"}}, exp{}, verifyConstLoopingVUs(I(3), 30*time.Second)},
 		{opts{cli: []string{"-u", "4", "--duration", "60s"}}, exp{}, verifyConstLoopingVUs(I(4), 1*time.Minute)},
 		{
-			opts{cli: []string{"--stage", "20s:10", "-s", "3m:5"}}, exp{},
+			opts{cli: []string{"--stage", "20s:10", "-s", "3m:5"}},
+			exp{},
 			verifyRampingVUs(null.NewInt(1, false), buildStages(20, 10, 180, 5)),
 		},
 		{
-			opts{cli: []string{"-s", "1m6s:5", "--vus", "10"}}, exp{},
+			opts{cli: []string{"-s", "1m6s:5", "--vus", "10"}},
+			exp{},
 			verifyRampingVUs(null.NewInt(10, true), buildStages(66, 5)),
 		},
 		{opts{cli: []string{"-u", "1", "-i", "6", "-d", "10s"}}, exp{}, func(t *testing.T, c Config) {
@@ -248,11 +250,13 @@ func getConfigConsolidationTestCases() []configConsolidationTestCase {
 		{opts{env: []string{"K6_VUS=5", "K6_ITERATIONS=15"}}, exp{}, verifySharedIters(I(5), I(15))},
 		{opts{env: []string{"K6_VUS=10", "K6_DURATION=20s"}}, exp{}, verifyConstLoopingVUs(I(10), 20*time.Second)},
 		{
-			opts{env: []string{"K6_STAGES=2m30s:11,1h1m:100"}}, exp{},
+			opts{env: []string{"K6_STAGES=2m30s:11,1h1m:100"}},
+			exp{},
 			verifyRampingVUs(null.NewInt(1, false), buildStages(150, 11, 3660, 100)),
 		},
 		{
-			opts{env: []string{"K6_STAGES=100s:100,0m30s:0", "K6_VUS=0"}}, exp{},
+			opts{env: []string{"K6_STAGES=100s:100,0m30s:0", "K6_VUS=0"}},
+			exp{},
 			verifyRampingVUs(null.NewInt(0, true), buildStages(100, 100, 30, 0)),
 		},
 		// Test if JSON configs work as expected
@@ -275,14 +279,16 @@ func getConfigConsolidationTestCases() []configConsolidationTestCase {
 				env: []string{"K6_DURATION=15s"},
 				cli: []string{"--stage", ""},
 			},
-			exp{logWarning: true}, verifyOneIterPerOneVU,
+			exp{logWarning: true},
+			verifyOneIterPerOneVU,
 		},
 		{
 			opts{
 				runner: &lib.Options{VUs: null.IntFrom(5), Duration: types.NullDurationFrom(50 * time.Second)},
 				cli:    []string{"--stage", "5s:5"},
 			},
-			exp{}, verifyRampingVUs(I(5), buildStages(5, 5)),
+			exp{},
+			verifyRampingVUs(I(5), buildStages(5, 5)),
 		},
 		{
 			opts{
@@ -323,7 +329,8 @@ func getConfigConsolidationTestCases() []configConsolidationTestCase {
 				env: []string{"K6_ITERATIONS=25"},
 				cli: []string{"--vus", "12"},
 			},
-			exp{}, verifySharedIters(I(12), I(25)),
+			exp{},
+			verifySharedIters(I(12), I(25)),
 		},
 
 		// TODO: test the externally controlled executor
@@ -375,6 +382,86 @@ func getConfigConsolidationTestCases() []configConsolidationTestCase {
 				assert.Equal(t, []string{"avg", "p(90)", "count"}, c.Options.SummaryTrendStats)
 			},
 		},
+		{opts{cli: []string{}}, exp{}, func(t *testing.T, c Config) {
+			assert.Equal(t, types.DNSConfig{
+				TTL:    null.NewString("5m", false),
+				Select: types.NullDNSSelect{DNSSelect: types.DNSrandom, Valid: false},
+				Policy: types.NullDNSPolicy{DNSPolicy: types.DNSpreferIPv4, Valid: false},
+			}, c.Options.DNS)
+		}},
+		{opts{env: []string{"K6_DNS=ttl=5,select=roundRobin"}}, exp{}, func(t *testing.T, c Config) {
+			assert.Equal(t, types.DNSConfig{
+				TTL:    null.StringFrom("5"),
+				Select: types.NullDNSSelect{DNSSelect: types.DNSroundRobin, Valid: true},
+				Policy: types.NullDNSPolicy{DNSPolicy: types.DNSpreferIPv4, Valid: false},
+			}, c.Options.DNS)
+		}},
+		{opts{env: []string{"K6_DNS=ttl=inf,select=random,policy=preferIPv6"}}, exp{}, func(t *testing.T, c Config) {
+			assert.Equal(t, types.DNSConfig{
+				TTL:    null.StringFrom("inf"),
+				Select: types.NullDNSSelect{DNSSelect: types.DNSrandom, Valid: true},
+				Policy: types.NullDNSPolicy{DNSPolicy: types.DNSpreferIPv6, Valid: true},
+			}, c.Options.DNS)
+		}},
+		// This is functionally invalid, but will error out in validation done in js.parseTTL().
+		{opts{cli: []string{"--dns", "ttl=-1"}}, exp{}, func(t *testing.T, c Config) {
+			assert.Equal(t, types.DNSConfig{
+				TTL:    null.StringFrom("-1"),
+				Select: types.NullDNSSelect{DNSSelect: types.DNSrandom, Valid: false},
+				Policy: types.NullDNSPolicy{DNSPolicy: types.DNSpreferIPv4, Valid: false},
+			}, c.Options.DNS)
+		}},
+		{opts{cli: []string{"--dns", "ttl=0,blah=nope"}}, exp{cliReadError: true}, nil},
+		{opts{cli: []string{"--dns", "ttl=0"}}, exp{}, func(t *testing.T, c Config) {
+			assert.Equal(t, types.DNSConfig{
+				TTL:    null.StringFrom("0"),
+				Select: types.NullDNSSelect{DNSSelect: types.DNSrandom, Valid: false},
+				Policy: types.NullDNSPolicy{DNSPolicy: types.DNSpreferIPv4, Valid: false},
+			}, c.Options.DNS)
+		}},
+		{opts{cli: []string{"--dns", "ttl=5s,select="}}, exp{cliReadError: true}, nil},
+		{
+			opts{fs: defaultConfig(`{"dns": {"ttl": "0", "select": "roundRobin", "policy": "onlyIPv4"}}`)},
+			exp{},
+			func(t *testing.T, c Config) {
+				assert.Equal(t, types.DNSConfig{
+					TTL:    null.StringFrom("0"),
+					Select: types.NullDNSSelect{DNSSelect: types.DNSroundRobin, Valid: true},
+					Policy: types.NullDNSPolicy{DNSPolicy: types.DNSonlyIPv4, Valid: true},
+				}, c.Options.DNS)
+			},
+		},
+		{
+			opts{
+				fs:  defaultConfig(`{"dns": {"ttl": "0"}}`),
+				env: []string{"K6_DNS=ttl=30,policy=any"},
+			},
+			exp{},
+			func(t *testing.T, c Config) {
+				assert.Equal(t, types.DNSConfig{
+					TTL:    null.StringFrom("30"),
+					Select: types.NullDNSSelect{DNSSelect: types.DNSrandom, Valid: false},
+					Policy: types.NullDNSPolicy{DNSPolicy: types.DNSany, Valid: true},
+				}, c.Options.DNS)
+			},
+		},
+		{
+			// CLI overrides all, falling back to env
+			opts{
+				fs:  defaultConfig(`{"dns": {"ttl": "60", "select": "first"}}`),
+				env: []string{"K6_DNS=ttl=30,select=random,policy=any"},
+				cli: []string{"--dns", "ttl=5"},
+			},
+			exp{},
+			func(t *testing.T, c Config) {
+				assert.Equal(t, types.DNSConfig{
+					TTL:    null.StringFrom("5"),
+					Select: types.NullDNSSelect{DNSSelect: types.DNSrandom, Valid: true},
+					Policy: types.NullDNSPolicy{DNSPolicy: types.DNSany, Valid: true},
+				}, c.Options.DNS)
+			},
+		},
+
 		// TODO: test for differences between flagsets
 		// TODO: more tests in general, especially ones not related to execution parameters...
 	}

cmd/options.go

@@ -93,6 +93,11 @@ func optionFlagSet() *pflag.FlagSet {
 	flags.StringSlice("tag", nil, "add a `tag` to be applied to all samples, as `[name]=[value]`")
 	flags.String("console-output", "", "redirects the console logging to the provided output file")
 	flags.Bool("discard-response-bodies", false, "Read but don't process or save HTTP response bodies")
+	flags.String("dns", types.DefaultDNSConfig().String(), "DNS resolver configuration. Possible ttl values are: 'inf' "+
+		"for a persistent cache, '0' to disable the cache,\nor a positive duration, e.g. '1s', '1m', etc. "+
+		"Milliseconds are assumed if no unit is provided.\n"+
+		"Possible select values to return a single IP are: 'first', 'random' or 'roundRobin'.\n"+
+		"Possible policy values are: 'preferIPv4', 'preferIPv6', 'onlyIPv4', 'onlyIPv6' or 'any'.\n")
 	return flags
 }
 
@@ -248,6 +253,14 @@ func getOptions(flags *pflag.FlagSet) (lib.Options, error) {
 		opts.ConsoleOutput = null.StringFrom(redirectConFile)
 	}
 
+	if dns, err := flags.GetString("dns"); err != nil {
+		return opts, err
+	} else if dns != "" {
+		if err := opts.DNS.UnmarshalText([]byte(dns)); err != nil {
+			return opts, err
+		}
+	}
+
 	return opts, nil
 }
 

core/local/local_test.go

@@ -24,6 +24,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"io/ioutil"
 	"net"
 	"net/url"
 	"reflect"
@@ -47,6 +48,7 @@ import (
 	"github.com/loadimpact/k6/lib/testutils"
 	"github.com/loadimpact/k6/lib/testutils/httpmultibin"
 	"github.com/loadimpact/k6/lib/testutils/minirunner"
+	"github.com/loadimpact/k6/lib/testutils/mockresolver"
 	"github.com/loadimpact/k6/lib/types"
 	"github.com/loadimpact/k6/loader"
 	"github.com/loadimpact/k6/stats"
@@ -974,6 +976,106 @@ func TestExecutionSchedulerIsRunning(t *testing.T) {
 	assert.NoError(t, <-err)
 }
 
+// TestDNSResolver checks the DNS resolution behavior at the ExecutionScheduler level.
+func TestDNSResolver(t *testing.T) {
+	tb := httpmultibin.NewHTTPMultiBin(t)
+	defer tb.Cleanup()
+	sr := tb.Replacer.Replace
+	script := sr(`
+		import http from "k6/http";
+		import { sleep } from "k6";
+
+		export let options = {
+			vus: 1,
+			iterations: 8,
+			noConnectionReuse: true,
+		}
+
+		export default function () {
+			const res = http.get("http://myhost:HTTPBIN_PORT/", { timeout: 50 });
+			sleep(0.7);  // somewhat uneven multiple of 0.5 to minimize races with asserts
+		}`)
+
+	t.Run("cache", func(t *testing.T) {
+		testCases := map[string]struct {
+			opts          lib.Options
+			expLogEntries int
+		}{
+			"default": { // IPs are cached for 5m
+				lib.Options{DNS: types.DefaultDNSConfig()}, 0,
+			},
+			"0": { // cache is disabled, every request does a DNS lookup
+				lib.Options{DNS: types.DNSConfig{
+					TTL:    null.StringFrom("0"),
+					Select: types.NullDNSSelect{DNSSelect: types.DNSfirst, Valid: true},
+					Policy: types.NullDNSPolicy{DNSPolicy: types.DNSpreferIPv4, Valid: false},
+				}}, 5,
+			},
+			"1000": { // cache IPs for 1s, check that unitless values are interpreted as ms
+				lib.Options{DNS: types.DNSConfig{
+					TTL:    null.StringFrom("1000"),
+					Select: types.NullDNSSelect{DNSSelect: types.DNSfirst, Valid: true},
+					Policy: types.NullDNSPolicy{DNSPolicy: types.DNSpreferIPv4, Valid: false},
+				}}, 4,
+			},
+			"3s": {
+				lib.Options{DNS: types.DNSConfig{
+					TTL:    null.StringFrom("3s"),
+					Select: types.NullDNSSelect{DNSSelect: types.DNSfirst, Valid: true},
+					Policy: types.NullDNSPolicy{DNSPolicy: types.DNSpreferIPv4, Valid: false},
+				}}, 3,
+			},
+		}
+
+		expErr := sr(`dial tcp 127.0.0.254:HTTPBIN_PORT: connect: connection refused`)
+		if runtime.GOOS == "windows" {
+			expErr = "context deadline exceeded"
+		}
+		for name, tc := range testCases {
+			tc := tc
+			t.Run(name, func(t *testing.T) {
+				logger := logrus.New()
+				logger.SetOutput(ioutil.Discard)
+				logHook := testutils.SimpleLogrusHook{HookedLevels: []logrus.Level{logrus.WarnLevel}}
+				logger.AddHook(&logHook)
+
+				runner, err := js.New(logger, &loader.SourceData{
+					URL: &url.URL{Path: "/script.js"}, Data: []byte(script),
+				}, nil, lib.RuntimeOptions{})
+				require.NoError(t, err)
+
+				mr := mockresolver.New(nil, net.LookupIP)
+				runner.ActualResolver = mr.LookupIPAll
+
+				ctx, cancel, execScheduler, samples := newTestExecutionScheduler(t, runner, logger, tc.opts)
+				defer cancel()
+
+				mr.Set("myhost", sr("HTTPBIN_IP"))
+				time.AfterFunc(1700*time.Millisecond, func() {
+					mr.Set("myhost", "127.0.0.254")
+				})
+				defer mr.Unset("myhost")
+
+				errCh := make(chan error, 1)
+				go func() { errCh <- execScheduler.Run(ctx, ctx, samples) }()
+
+				select {
+				case err := <-errCh:
+					require.NoError(t, err)
+					entries := logHook.Drain()
+					require.Len(t, entries, tc.expLogEntries)
+					for _, entry := range entries {
+						require.IsType(t, &url.Error{}, entry.Data["error"])
+						assert.EqualError(t, entry.Data["error"].(*url.Error).Err, expErr)
+					}
+				case <-time.After(10 * time.Second):
+					t.Fatal("timed out")
+				}
+			})
+		}
+	})
+}
+
 func TestRealTimeAndSetupTeardownMetrics(t *testing.T) {
 	if runtime.GOOS == "windows" {
 		t.Skip()
@@ -1100,7 +1202,10 @@ func TestRealTimeAndSetupTeardownMetrics(t *testing.T) {
 	getDummyTrail := func(group string, emitIterations bool, addExpTags ...string) stats.SampleContainer {
 		expTags := []string{"group", group}
 		expTags = append(expTags, addExpTags...)
-		return netext.NewDialer(net.Dialer{}).GetTrail(time.Now(), time.Now(),
+		return netext.NewDialer(
+			net.Dialer{},
+			netext.NewResolver(net.LookupIP, 0, types.DNSfirst, types.DNSpreferIPv4),
+		).GetTrail(time.Now(), time.Now(),
 			true, emitIterations, getTags(expTags...))
 	}
 

go.mod

@@ -66,7 +66,6 @@ require (
 	github.com/urfave/negroni v0.3.1-0.20180130044549-22c5532ea862
 	github.com/valyala/bytebufferpool v1.0.0 // indirect
 	github.com/valyala/fasttemplate v0.0.0-20170224212429-dcecefd839c4 // indirect
-	github.com/viki-org/dnscache v0.0.0-20130720023526-c70c1f23c5d8
 	github.com/zyedidia/highlight v0.0.0-20170330143449-201131ce5cf5
 	golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2
 	golang.org/x/net v0.0.0-20200520004742-59133d7f0dd7

go.sum

@@ -153,8 +153,6 @@ github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6Kllzaw
 github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc=
 github.com/valyala/fasttemplate v0.0.0-20170224212429-dcecefd839c4 h1:gKMu1Bf6QINDnvyZuTaACm9ofY+PRh+5vFz4oxBZeF8=
 github.com/valyala/fasttemplate v0.0.0-20170224212429-dcecefd839c4/go.mod h1:50wTf68f99/Zt14pr046Tgt3Lp2vLyFZKzbFXTOabXw=
-github.com/viki-org/dnscache v0.0.0-20130720023526-c70c1f23c5d8 h1:EVObHAr8DqpoJCVv6KYTle8FEImKhtkfcZetNqxDoJQ=
-github.com/viki-org/dnscache v0.0.0-20130720023526-c70c1f23c5d8/go.mod h1:dniwbG03GafCjFohMDmz6Zc6oCuiqgH6tGNyXTkHzXE=
 github.com/zyedidia/highlight v0.0.0-20170330143449-201131ce5cf5 h1:Zs6mpwXvlqpF9zHl5XaN0p5V4J9XvP+WBuiuXyIgqvc=
 github.com/zyedidia/highlight v0.0.0-20170330143449-201131ce5cf5/go.mod h1:c1r+Ob9tUTPB0FKWO1+x+Hsc/zNa45WdGq7Y38Ybip0=
 golang.org/x/crypto v0.0.0-20180308185624-c7dcf104e3a7 h1:c9Tyi4qyEZwEJ1+Zm6Fcqf+68wmUdMzfXYTp3s8Nzg8=

js/bundle.go

@@ -249,7 +249,7 @@ func (b *Bundle) Instantiate(logger logrus.FieldLogger, vuID int64) (bi *BundleI
 	}
 
 	// Grab any exported functions that could be executed. These were
-	// already pre-validated in NewBundle(), just get them here.
+	// already pre-validated in cmd.validateScenarioConfig(), just get them here.
 	exports := rt.Get("exports").ToObject(rt)
 	for k := range b.exports {
 		fn, _ := goja.AssertFunction(exports.Get(k))

js/initcontext_test.go

@@ -43,6 +43,7 @@ import (
 	"github.com/loadimpact/k6/lib/consts"
 	"github.com/loadimpact/k6/lib/netext"
 	"github.com/loadimpact/k6/lib/testutils"
+	"github.com/loadimpact/k6/lib/types"
 	"github.com/loadimpact/k6/stats"
 )
 
@@ -388,11 +389,14 @@ func TestRequestWithBinaryFile(t *testing.T) {
 		Logger:  logger,
 		Group:   root,
 		Transport: &http.Transport{
-			DialContext: (netext.NewDialer(net.Dialer{
-				Timeout:   10 * time.Second,
-				KeepAlive: 60 * time.Second,
-				DualStack: true,
-			})).DialContext,
+			DialContext: (netext.NewDialer(
+				net.Dialer{
+					Timeout:   10 * time.Second,
+					KeepAlive: 60 * time.Second,
+					DualStack: true,
+				},
+				netext.NewResolver(net.LookupIP, 0, types.DNSfirst, types.DNSpreferIPv4),
+			)).DialContext,
 		},
 		BPool:   bpool.NewBufferPool(1),
 		Samples: make(chan stats.SampleContainer, 500),