forked from cisco-app-networking/networkservicemesh
-
Notifications
You must be signed in to change notification settings - Fork 0
/
context.go
113 lines (99 loc) · 3.57 KB
/
context.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
package endpoint
import (
"context"
connectionMonitor "github.com/networkservicemesh/networkservicemesh/sdk/monitor/connectionmonitor"
"github.com/sirupsen/logrus"
"github.com/networkservicemesh/networkservicemesh/controlplane/api/connection"
"github.com/networkservicemesh/networkservicemesh/controlplane/api/networkservice"
)
type contextKeyType string
const (
clientConnectionKey contextKeyType = "ClientConnection"
monitorServerKey contextKeyType = "MonitorServer"
nextKey contextKeyType = "Next"
logKey contextKeyType = "Log"
)
// WithClientConnection -
// Wraps 'parent' in a new Context that has the ClientConnection
// provided in:
// connection *connection.Connection
// using Context.Value(...) and returns the result.
// Note: any previously existing ClientConnection will be overwritten.
//
// Recommended use: in any Request or Close call that creates a ClientConnection, call:
// ctx = WithClientConnection(ctx)
// to ensure that the ctx has a ClientConnection
// In any Request or Close call that consumes a ClientConnection, call:
// connectionMap := ClientConnection(ctx)
// to retrieve the ClientConnection from the context.Context
func WithClientConnection(parent context.Context, connection *connection.Connection) context.Context {
if parent == nil {
parent = context.Background()
}
return context.WithValue(parent, clientConnectionKey, connection)
}
// ClientConnection -
// Returns a ClientConnection from:
// ctx context.Context
// If any is present, otherwise nil
func ClientConnection(ctx context.Context) *connection.Connection {
conn := ctx.Value(clientConnectionKey)
if conn == nil {
return nil
}
return conn.(*connection.Connection)
}
// withNext -
// Wraps 'parent' in a new Context that has the Next networkservice.NetworkServiceServer to be called in the chain
// Should only be set in CompositeEndpoint.Request/Close
func withNext(parent context.Context, next networkservice.NetworkServiceServer) context.Context {
if parent == nil {
parent = context.TODO()
}
return context.WithValue(parent, nextKey, next)
}
// Next -
// Returns the Next networkservice.NetworkServiceServer to be called in the chain from the context.Context
func Next(ctx context.Context) networkservice.NetworkServiceServer {
if rv, ok := ctx.Value(nextKey).(networkservice.NetworkServiceServer); ok {
return rv
}
return nil
}
// withLog -
// Provides a FieldLogger in context
func withLog(parent context.Context, log logrus.FieldLogger) context.Context {
if parent == nil {
parent = context.TODO()
}
return context.WithValue(parent, logKey, log)
}
// Log - return FieldLogger from context
func Log(ctx context.Context) logrus.FieldLogger {
if rv, ok := ctx.Value(logKey).(logrus.FieldLogger); ok {
return rv
}
return logrus.New()
}
// WithMonitorServer -
// Wraps 'parent' in a new Context that has the local connection Monitor
// using Context.Value(...) and returns the result.
// Note: any previously existing MonitorServer will be overwritten.
//
func WithMonitorServer(parent context.Context, monitorServer connectionMonitor.MonitorServer) context.Context {
if parent == nil {
parent = context.Background()
}
return context.WithValue(parent, monitorServerKey, monitorServer)
}
// MonitorServer -
// Returns a MonitorServer from:
// ctx context.Context
// If any is present, otherwise nil
func MonitorServer(ctx context.Context) connectionMonitor.MonitorServer {
value := ctx.Value(monitorServerKey)
if value == nil {
return nil
}
return value.(connectionMonitor.MonitorServer)
}