66 "bytes"
77 "fmt"
88 "io/ioutil"
9+ "io"
910 "math/rand"
1011 "net/http"
1112 "os"
@@ -67,6 +68,11 @@ func TestBasicAuthConfig(t *testing.T) {
6768 t .Run ("WithPerSecondRedisAuthWithLocalCache" , testBasicConfigAuth ("18093" , "true" , "1000" ))
6869}
6970
71+ func TestBasicReloadConfig (t * testing.T ) {
72+ t .Run ("BasicWithoutWatchRoot" , testBasicConfigWithoutWatchRoot ("8095" , "false" , "0" ))
73+ t .Run ("ReloadWithoutWatchRoot" , testBasicConfigReload ("8097" , "false" , "0" , "false" ))
74+ }
75+
7076func testBasicConfigAuthTLS (grpcPort , perSecond string , local_cache_size string ) func (* testing.T ) {
7177 os .Setenv ("REDIS_PERSECOND_URL" , "localhost:16382" )
7278 os .Setenv ("REDIS_URL" , "localhost:16381" )
@@ -97,6 +103,28 @@ func testBasicConfigAuth(grpcPort, perSecond string, local_cache_size string) fu
97103 return testBasicBaseConfig (grpcPort , perSecond , local_cache_size )
98104}
99105
106+ func testBasicConfigWithoutWatchRoot (grpcPort , perSecond string , local_cache_size string ) func (* testing.T ) {
107+ os .Setenv ("REDIS_PERSECOND_URL" , "localhost:6380" )
108+ os .Setenv ("REDIS_URL" , "localhost:6379" )
109+ os .Setenv ("REDIS_AUTH" , "" )
110+ os .Setenv ("REDIS_TLS" , "false" )
111+ os .Setenv ("REDIS_PERSECOND_AUTH" , "" )
112+ os .Setenv ("REDIS_PERSECOND_TLS" , "false" )
113+ os .Setenv ("RUNTIME_WATCH_ROOT" , "false" )
114+ return testBasicBaseConfig (grpcPort , perSecond , local_cache_size )
115+ }
116+
117+ func testBasicConfigReload (grpcPort , perSecond string , local_cache_size , runtimeWatchRoot string ) func (* testing.T ) {
118+ os .Setenv ("REDIS_PERSECOND_URL" , "localhost:6380" )
119+ os .Setenv ("REDIS_URL" , "localhost:6379" )
120+ os .Setenv ("REDIS_AUTH" , "" )
121+ os .Setenv ("REDIS_TLS" , "false" )
122+ os .Setenv ("REDIS_PERSECOND_AUTH" , "" )
123+ os .Setenv ("REDIS_PERSECOND_TLS" , "false" )
124+ os .Setenv ("RUNTIME_WATCH_ROOT" , runtimeWatchRoot )
125+ return testConfigReload (grpcPort , perSecond , local_cache_size )
126+ }
127+
100128func getCacheKey (cacheKey string , enableLocalCache bool ) string {
101129 if enableLocalCache {
102130 return cacheKey + "_local"
@@ -456,3 +484,105 @@ func TestBasicConfigLegacy(t *testing.T) {
456484 assert .NoError (err )
457485 }
458486}
487+
488+ func testConfigReload (grpcPort , perSecond string , local_cache_size string ) func (* testing.T ) {
489+ return func (t * testing.T ) {
490+ os .Setenv ("REDIS_PERSECOND" , perSecond )
491+ os .Setenv ("PORT" , "8082" )
492+ os .Setenv ("GRPC_PORT" , grpcPort )
493+ os .Setenv ("DEBUG_PORT" , "8084" )
494+ os .Setenv ("RUNTIME_ROOT" , "runtime/current" )
495+ os .Setenv ("RUNTIME_SUBDIRECTORY" , "ratelimit" )
496+ os .Setenv ("REDIS_PERSECOND_SOCKET_TYPE" , "tcp" )
497+ os .Setenv ("REDIS_SOCKET_TYPE" , "tcp" )
498+ os .Setenv ("LOCAL_CACHE_SIZE_IN_BYTES" , local_cache_size )
499+ os .Setenv ("USE_STATSD" , "false" )
500+
501+ local_cache_size_val , _ := strconv .Atoi (local_cache_size )
502+ enable_local_cache := local_cache_size_val > 0
503+ runner := runner .NewRunner ()
504+
505+ go func () {
506+ runner .Run ()
507+ }()
508+
509+ // HACK: Wait for the server to come up. Make a hook that we can wait on.
510+ time .Sleep (1 * time .Second )
511+
512+ assert := assert .New (t )
513+ conn , err := grpc .Dial (fmt .Sprintf ("localhost:%s" , grpcPort ), grpc .WithInsecure ())
514+ assert .NoError (err )
515+ defer conn .Close ()
516+ c := pb .NewRateLimitServiceClient (conn )
517+
518+ response , err := c .ShouldRateLimit (
519+ context .Background (),
520+ common .NewRateLimitRequest ("reload" , [][][2 ]string {{{getCacheKey ("block" , enable_local_cache ), "foo" }}}, 1 ))
521+ assert .Equal (
522+ & pb.RateLimitResponse {
523+ OverallCode : pb .RateLimitResponse_OK ,
524+ Statuses : []* pb.RateLimitResponse_DescriptorStatus {{Code : pb .RateLimitResponse_OK }}},
525+ response )
526+ assert .NoError (err )
527+
528+ runner .GetStatsStore ().Flush ()
529+ loadCount1 := runner .GetStatsStore ().NewCounter ("ratelimit.service.config_load_success" ).Value ()
530+
531+ // Copy a new file to config folder to test config reload functionality
532+ in , err := os .Open ("runtime/current/ratelimit/reload.yaml" )
533+ if err != nil {
534+ panic (err )
535+ }
536+ defer in .Close ()
537+ out , err := os .Create ("runtime/current/ratelimit/config/reload.yaml" )
538+ if err != nil {
539+ panic (err )
540+ }
541+ defer out .Close ()
542+ _ , err = io .Copy (out , in )
543+ if err != nil {
544+ panic (err )
545+ }
546+ err = out .Close ()
547+ if err != nil {
548+ panic (err )
549+ }
550+
551+ // Need to wait for config reload to take place and new descriptors to be loaded.
552+ // Shouldn't take more than 5 seconds but wait 120 at most just to be safe.
553+ wait := 120
554+ reloaded := false
555+ loadCount2 := uint64 (0 )
556+
557+ for i := 0 ; i < wait ; i ++ {
558+ time .Sleep (1 * time .Second )
559+ runner .GetStatsStore ().Flush ()
560+ loadCount2 = runner .GetStatsStore ().NewCounter ("ratelimit.service.config_load_success" ).Value ()
561+
562+ // Check that successful loads count has increased before continuing.
563+ if loadCount2 > loadCount1 {
564+ reloaded = true
565+ break
566+ }
567+ }
568+
569+ assert .True (reloaded )
570+ assert .Greater (loadCount2 , loadCount1 )
571+
572+ response , err = c .ShouldRateLimit (
573+ context .Background (),
574+ common .NewRateLimitRequest ("reload" , [][][2 ]string {{{getCacheKey ("key1" , enable_local_cache ), "foo" }}}, 1 ))
575+ assert .Equal (
576+ & pb.RateLimitResponse {
577+ OverallCode : pb .RateLimitResponse_OK ,
578+ Statuses : []* pb.RateLimitResponse_DescriptorStatus {
579+ newDescriptorStatus (pb .RateLimitResponse_OK , 50 , pb .RateLimitResponse_RateLimit_SECOND , 49 )}},
580+ response )
581+ assert .NoError (err )
582+
583+ err = os .Remove ("runtime/current/ratelimit/config/reload.yaml" )
584+ if err != nil {
585+ panic (err )
586+ }
587+ }
588+ }
0 commit comments