go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/server/redisconn/module.go (about)

     1  // Copyright 2020 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package redisconn
    16  
    17  import (
    18  	"context"
    19  	"flag"
    20  	"fmt"
    21  
    22  	"go.chromium.org/luci/common/logging"
    23  	"go.chromium.org/luci/common/tsmon"
    24  
    25  	"go.chromium.org/luci/server/caching"
    26  	"go.chromium.org/luci/server/module"
    27  	"go.chromium.org/luci/server/redisconn/adminpb"
    28  )
    29  
    30  // ModuleName can be used to refer to this module when declaring dependencies.
    31  var ModuleName = module.RegisterName("go.chromium.org/luci/server/redisconn")
    32  
    33  // ModuleOptions contain configuration of the Redis server module.
    34  type ModuleOptions struct {
    35  	RedisAddr string // Redis server to connect to as "host:port"
    36  	RedisDB   int    // index of a logical Redis DB to use by default
    37  }
    38  
    39  // Register registers the command line flags.
    40  func (o *ModuleOptions) Register(f *flag.FlagSet) {
    41  	f.StringVar(
    42  		&o.RedisAddr,
    43  		"redis-addr",
    44  		o.RedisAddr,
    45  		`Redis server to connect to as "host:port" (optional, Redis calls won't work if not given)`,
    46  	)
    47  	f.IntVar(
    48  		&o.RedisDB,
    49  		"redis-db",
    50  		o.RedisDB,
    51  		fmt.Sprintf("Index of a logical Redis DB to use by default (default is %d)", o.RedisDB),
    52  	)
    53  }
    54  
    55  // NewModule returns a server module that adds a Redis connection pool to the
    56  // global server context and installs Redis as the default caching.BlobCache
    57  // implementation.
    58  //
    59  // The Redis connection pool can be used through redisconn.Get(ctx).
    60  //
    61  // Does nothing if RedisAddr options is unset. In this case redisconn.Get(ctx)
    62  // returns ErrNotConfigured.
    63  func NewModule(opts *ModuleOptions) module.Module {
    64  	if opts == nil {
    65  		opts = &ModuleOptions{}
    66  	}
    67  	return &redisModule{opts: opts}
    68  }
    69  
    70  // NewModuleFromFlags is a variant of NewModule that initializes options through
    71  // command line flags.
    72  //
    73  // Calling this function registers flags in flag.CommandLine. They are usually
    74  // parsed in server.Main(...).
    75  func NewModuleFromFlags() module.Module {
    76  	opts := &ModuleOptions{}
    77  	opts.Register(flag.CommandLine)
    78  	return NewModule(opts)
    79  }
    80  
    81  // redisModule implements module.Module.
    82  type redisModule struct {
    83  	opts *ModuleOptions
    84  }
    85  
    86  // Name is part of module.Module interface.
    87  func (*redisModule) Name() module.Name {
    88  	return ModuleName
    89  }
    90  
    91  // Dependencies is part of module.Module interface.
    92  func (*redisModule) Dependencies() []module.Dependency {
    93  	return nil
    94  }
    95  
    96  // Initialize is part of module.Module interface.
    97  func (m *redisModule) Initialize(ctx context.Context, host module.Host, opts module.HostOptions) (context.Context, error) {
    98  	if m.opts.RedisAddr == "" {
    99  		return ctx, nil
   100  	}
   101  
   102  	pool := NewPool(m.opts.RedisAddr, m.opts.RedisDB)
   103  	ctx = UsePool(ctx, pool)
   104  
   105  	// Use Redis as caching.BlobCache provider.
   106  	ctx = caching.WithGlobalCache(ctx, func(namespace string) caching.BlobCache {
   107  		return &redisBlobCache{Prefix: fmt.Sprintf("luci.blobcache.%s:", namespace)}
   108  	})
   109  
   110  	// Close all connections when exiting gracefully.
   111  	host.RegisterCleanup(func(ctx context.Context) {
   112  		if err := pool.Close(); err != nil {
   113  			logging.Warningf(ctx, "Failed to close Redis pool - %s", err)
   114  		}
   115  	})
   116  
   117  	// Populate pool metrics on tsmon flush.
   118  	tsmon.RegisterCallbackIn(ctx, func(ctx context.Context) {
   119  		ReportStats(ctx, pool, "default")
   120  	})
   121  
   122  	// Expose an admin API that can be used to e.g. flush Redis DB. This is
   123  	// especially useful on GAE where reaching Redis otherwise requires launching
   124  	// a VM to get into the private network.
   125  	adminpb.RegisterAdminServer(host, &adminServer{pool: pool})
   126  
   127  	return ctx, nil
   128  }