@@ -25,8 +25,13 @@ import (
2525 "strconv"
2626 "time"
2727
28+ "github.com/cortexlabs/cortex/pkg/lib/aws"
29+ "github.com/cortexlabs/cortex/pkg/lib/errors"
2830 "github.com/cortexlabs/cortex/pkg/lib/logging"
31+ "github.com/cortexlabs/cortex/pkg/lib/telemetry"
2932 "github.com/cortexlabs/cortex/pkg/proxy"
33+ "github.com/cortexlabs/cortex/pkg/types/clusterconfig"
34+ "github.com/cortexlabs/cortex/pkg/types/userconfig"
3035 "go.uber.org/zap"
3136)
3237
@@ -35,20 +40,45 @@ const (
3540 _requestSampleInterval = 1 * time .Second
3641)
3742
43+ var (
44+ proxyLogger = logging .GetLogger ()
45+ )
46+
47+ func Exit (err error , wrapStrs ... string ) {
48+ for _ , str := range wrapStrs {
49+ err = errors .Wrap (err , str )
50+ }
51+
52+ if err != nil && ! errors .IsNoTelemetry (err ) {
53+ telemetry .Error (err )
54+ }
55+
56+ if err != nil && ! errors .IsNoPrint (err ) {
57+ proxyLogger .Error (err )
58+ }
59+
60+ telemetry .Close ()
61+
62+ os .Exit (1 )
63+ }
64+
3865func main () {
3966 var (
4067 port int
4168 metricsPort int
4269 userContainerPort int
4370 maxConcurrency int
4471 maxQueueLength int
72+ clusterConfigPath string
4573 )
4674
47- flag .IntVar (& port , "port" , 8000 , "port where the proxy will be served" )
48- flag .IntVar (& metricsPort , "metrics-port" , 8001 , "port where the proxy will be served " )
49- flag .IntVar (& userContainerPort , "user-port" , 8080 , "port where the proxy will redirect to the traffic to" )
75+ flag .IntVar (& port , "port" , 8888 , "port where the proxy is served" )
76+ flag .IntVar (& metricsPort , "metrics-port" , 15000 , "metrics port for prometheus " )
77+ flag .IntVar (& userContainerPort , "user-port" , 8080 , "port where the proxy redirects to the traffic to" )
5078 flag .IntVar (& maxConcurrency , "max-concurrency" , 0 , "max concurrency allowed for user container" )
5179 flag .IntVar (& maxQueueLength , "max-queue-length" , 0 , "max request queue length for user container" )
80+ flag .StringVar (& clusterConfigPath , "cluster-config" , "" , "cluster config path" )
81+
5282 flag .Parse ()
5383
5484 log := logging .GetLogger ()
@@ -58,12 +88,44 @@ func main() {
5888
5989 switch {
6090 case maxConcurrency == 0 :
61- log .Fatal ("-- max-concurrency flag is required" )
91+ log .Fatal ("-max-concurrency flag is required" )
6292 case maxQueueLength == 0 :
63- maxQueueLength = maxConcurrency * 10
93+ log .Fatal ("-max-queue-length flag is required" )
94+ case clusterConfigPath == "" :
95+ log .Fatal ("-cluster-config flag is required" )
96+ }
97+
98+ clusterConfig , err := clusterconfig .NewForFile (clusterConfigPath )
99+ if err != nil {
100+ Exit (err )
101+ }
102+
103+ awsClient , err := aws .NewForRegion (clusterConfig .Region )
104+ if err != nil {
105+ Exit (err )
106+ }
107+
108+ _ , userID , err := awsClient .CheckCredentials ()
109+ if err != nil {
110+ Exit (err )
111+ }
112+
113+ err = telemetry .Init (telemetry.Config {
114+ Enabled : clusterConfig .Telemetry ,
115+ UserID : userID ,
116+ Properties : map [string ]string {
117+ "kind" : userconfig .RealtimeAPIKind .String (),
118+ "image_type" : "proxy" ,
119+ },
120+ Environment : "api" ,
121+ LogErrors : true ,
122+ BackoffMode : telemetry .BackoffDuplicateMessages ,
123+ })
124+ if err != nil {
125+ Exit (err )
64126 }
65127
66- target := "http://127.0.0.1:" + strconv .Itoa (port )
128+ target := "http://127.0.0.1:" + strconv .Itoa (userContainerPort )
67129 httpProxy := proxy .NewReverseProxy (target , maxQueueLength , maxQueueLength )
68130
69131 requestCounterStats := & proxy.RequestStats {}
@@ -101,7 +163,7 @@ func main() {
101163
102164 servers := map [string ]* http.Server {
103165 "proxy" : {
104- Addr : ":" + strconv .Itoa (userContainerPort ),
166+ Addr : ":" + strconv .Itoa (port ),
105167 Handler : proxy .Handler (breaker , httpProxy ),
106168 },
107169 "metrics" : {
@@ -123,7 +185,7 @@ func main() {
123185
124186 select {
125187 case err := <- errCh :
126- log . Fatal ( "failed to start proxy server" , zap . Error ( err ))
188+ Exit ( errors . Wrap ( err , "failed to start proxy server" ))
127189 case <- sigint :
128190 // We received an interrupt signal, shut down.
129191 log .Info ("Received TERM signal, handling a graceful shutdown..." )
@@ -133,8 +195,10 @@ func main() {
133195 if err := server .Shutdown (context .Background ()); err != nil {
134196 // Error from closing listeners, or context timeout:
135197 log .Warn ("HTTP server Shutdown Error" , zap .Error (err ))
198+ telemetry .Error (errors .Wrap (err , "HTTP server Shutdown Error" ))
136199 }
137200 }
138201 log .Info ("Shutdown complete, exiting..." )
202+ telemetry .Close ()
139203 }
140204}
0 commit comments