[m3admin-client] Set zone header based on cluster spec (#289)

integration/harness/m3admin.go

@@ -42,6 +42,7 @@ func (h *Harness) NewPlacementClient(cluster *myspec.M3DBCluster) (placement.Cli
 	env := k8sops.DefaultM3ClusterEnvironmentName(cluster)
 	adminCl := m3admin.NewClient(
 		m3admin.WithEnvironment(env),
+		m3admin.WithZone(cluster.Spec.Zone),
 	)
 	url := fmt.Sprintf(proxyBaseURLFmt, h.Namespace, svc.Name, "7201")
 	h.Logger.Sugar().Infof("calling url '%s' with env '%s'", url, env)

pkg/controller/m3admin_client.go

@@ -109,7 +109,7 @@ func newMultiAdminClient(adminOpts []m3admin.Option, logger *zap.Logger) *multiA
 
 func (m *multiAdminClient) adminClientForCluster(cluster *myspec.M3DBCluster) m3admin.Client {
 	env := k8sops.DefaultM3ClusterEnvironmentName(cluster)
-	opts := append(m.adminOpts, m3admin.WithEnvironment(env))
+	opts := append(m.adminOpts, m3admin.WithEnvironment(env), m3admin.WithZone(cluster.Spec.Zone))
 	return m.adminClientFn(opts...)
 }
 

pkg/m3admin/client.go

@@ -38,6 +38,8 @@ import (
 
 const (
 	m3EnvironmentHeader = "Cluster-Environment-Name"
+	// m3HeaderClusterZoneName is the header used to specify the zone name.
+	m3HeaderClusterZoneName = "Cluster-Zone-Name"
 )
 
 var (
@@ -61,6 +63,7 @@ type client struct {
 	client      *retryhttp.Client
 	logger      *zap.Logger
 	environment string
+	zone        string
 }
 
 type nullLogger struct{}
@@ -78,6 +81,7 @@ func NewClient(clientOpts ...Option) Client {
 		client:      opts.client,
 		logger:      opts.logger,
 		environment: opts.environment,
+		zone:        opts.zone,
 	}
 
 	if client.client == nil {
@@ -133,6 +137,9 @@ func (c *client) DoHTTPRequest(
 	if c.environment != "" {
 		request.Header.Add(m3EnvironmentHeader, c.environment)
 	}
+	if c.zone != "" {
+		request.Header.Add(m3HeaderClusterZoneName, c.zone)
+	}
 
 	if l.Core().Enabled(zapcore.DebugLevel) {
 		dump, err := httputil.DumpRequest(request.Request, true)

pkg/m3admin/client_test.go

@@ -92,6 +92,10 @@ func TestClient_DoHTTPRequest_Header(t *testing.T) {
 			w.WriteHeader(500)
 			return
 		}
+		if r.Header.Get(m3HeaderClusterZoneName) != "bar-zone" {
+			w.WriteHeader(500)
+			return
+		}
 		w.Write([]byte("hello"))
 	}))
 	defer s.Close()
@@ -102,11 +106,15 @@ func TestClient_DoHTTPRequest_Header(t *testing.T) {
 		return data
 	}
 
-	cl := newTestClient(WithEnvironment("fooz-env"))
+	cl := newTestClient(WithEnvironment("fooz-env"), WithZone("bar-zone"))
 	_, err := cl.DoHTTPRequest("GET", s.URL, nil)
 	assert.Error(t, err)
 
-	cl = newTestClient(WithEnvironment("foo-env"))
+	cl = newTestClient(WithEnvironment("foo-env"), WithZone("barz-zone"))
+	_, err = cl.DoHTTPRequest("GET", s.URL, nil) //nolint:bodyclose
+	assert.Error(t, err)
+
+	cl = newTestClient(WithEnvironment("foo-env"), WithZone("bar-zone"))
 	resp, err := cl.DoHTTPRequest("GET", s.URL, nil)
 	require.NoError(t, err)
 	assert.Equal(t, resp.StatusCode, http.StatusOK)

pkg/m3admin/options.go

@@ -40,6 +40,7 @@ type options struct {
 	logger      *zap.Logger
 	client      *retryhttp.Client
 	environment string
+	zone        string
 }
 
 // WithLogger configures a logger for the client. If not set a noop logger will
@@ -67,6 +68,13 @@ func WithEnvironment(e string) Option {
 	})
 }
 
+// WithZone controls the Cluster-Zone-Name header to m3coordinator.
+func WithZone(z string) Option {
+	return optionFn(func(o *options) {
+		o.zone = z
+	})
+}
+
 // RequestOption defines a per-request option.
 type RequestOption interface {
 	execute(*reqOptions)