Add ttrpc namespace support

crosbymichael · crosbymichael · commit fa6a9f063586 · 2019-06-04T15:21:10.000Z
Signed-off-by: Michael Crosby &lt;crosbymichael@gmail.com&gt;
diff --git a/namespaces/context.go b/namespaces/context.go
@@ -36,10 +36,9 @@ type namespaceKey struct{}
 // WithNamespace sets a given namespace on the context
 func WithNamespace(ctx context.Context, namespace string) context.Context {
 	ctx = context.WithValue(ctx, namespaceKey{}, namespace) // set our key for namespace
-
-	// also store on the grpc headers so it gets picked up by any clients that
+	// also store on the grpc and ttrpc headers so it gets picked up by any clients that
 	// are using this.
-	return withGRPCNamespaceHeader(ctx, namespace)
+	return withTTRPCNamespaceHeader(withGRPCNamespaceHeader(ctx, namespace), namespace)
 }
 
 // NamespaceFromEnv uses the namespace defined in CONTAINERD_NAMESPACE or
@@ -58,9 +57,10 @@ func NamespaceFromEnv(ctx context.Context) context.Context {
 func Namespace(ctx context.Context) (string, bool) {
 	namespace, ok := ctx.Value(namespaceKey{}).(string)
 	if !ok {
-		return fromGRPCHeader(ctx)
+		if namespace, ok = fromGRPCHeader(ctx); !ok {
+			return fromTTRPCHeader(ctx)
+		}
 	}
-
 	return namespace, ok
 }
 
@@ -70,10 +70,8 @@ func NamespaceRequired(ctx context.Context) (string, error) {
 	if !ok || namespace == "" {
 		return "", errors.Wrapf(errdefs.ErrFailedPrecondition, "namespace is required")
 	}
-
 	if err := Validate(namespace); err != nil {
 		return "", errors.Wrap(err, "namespace validation")
 	}
-
 	return namespace, nil
 }
diff --git a/namespaces/ttrpc.go b/namespaces/ttrpc.go
@@ -0,0 +1,41 @@
+/*
+   Copyright The containerd Authors.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+package namespaces
+
+import (
+	"context"
+
+	"github.com/containerd/ttrpc"
+)
+
+const (
+	// TTRPCHeader defines the header name for specifying a containerd namespace
+	TTRPCHeader = "containerd-namespace-ttrpc"
+)
+
+func withTTRPCNamespaceHeader(ctx context.Context, namespace string) context.Context {
+	md, ok := ttrpc.GetMetadata(ctx)
+	if !ok {
+		md = ttrpc.Metadata{}
+	}
+	md.Set(TTRPCHeader, namespace)
+	return ttrpc.WithMetadata(ctx, md)
+}
+
+func fromTTRPCHeader(ctx context.Context) (string, bool) {
+	return ttrpc.GetMetadataValue(ctx, TTRPCHeader)
+}

Original file line number	Diff line number	Diff line change
`@@ -36,10 +36,9 @@ type namespaceKey struct{}`
`36`	`36`	`// WithNamespace sets a given namespace on the context`
`37`	`37`	`func WithNamespace(ctx context.Context, namespace string) context.Context {`
`38`	`38`	`ctx = context.WithValue(ctx, namespaceKey{}, namespace) // set our key for namespace`
`39`		`-`
`40`		`- // also store on the grpc headers so it gets picked up by any clients that`
	`39`	`+ // also store on the grpc and ttrpc headers so it gets picked up by any clients that`
`41`	`40`	`// are using this.`
`42`		`- return withGRPCNamespaceHeader(ctx, namespace)`
	`41`	`+ return withTTRPCNamespaceHeader(withGRPCNamespaceHeader(ctx, namespace), namespace)`
`43`	`42`	`}`
`44`	`43`
`45`	`44`	`// NamespaceFromEnv uses the namespace defined in CONTAINERD_NAMESPACE or`
`@@ -58,9 +57,10 @@ func NamespaceFromEnv(ctx context.Context) context.Context {`
`58`	`57`	`func Namespace(ctx context.Context) (string, bool) {`
`59`	`58`	`namespace, ok := ctx.Value(namespaceKey{}).(string)`
`60`	`59`	`if !ok {`
`61`		`- return fromGRPCHeader(ctx)`
	`60`	`+ if namespace, ok = fromGRPCHeader(ctx); !ok {`
	`61`	`+ return fromTTRPCHeader(ctx)`
	`62`	`+ }`
`62`	`63`	`}`
`63`		`-`
`64`	`64`	`return namespace, ok`
`65`	`65`	`}`
`66`	`66`
`@@ -70,10 +70,8 @@ func NamespaceRequired(ctx context.Context) (string, error) {`
`70`	`70`	`if !ok \|\| namespace == "" {`
`71`	`71`	`return "", errors.Wrapf(errdefs.ErrFailedPrecondition, "namespace is required")`
`72`	`72`	`}`
`73`		`-`
`74`	`73`	`if err := Validate(namespace); err != nil {`
`75`	`74`	`return "", errors.Wrap(err, "namespace validation")`
`76`	`75`	`}`
`77`		`-`
`78`	`76`	`return namespace, nil`
`79`	`77`	`}`