go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/casviewer/client.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 casviewer
    16  
    17  import (
    18  	"context"
    19  	"sync"
    20  
    21  	"github.com/bazelbuild/remote-apis-sdks/go/pkg/client"
    22  	"google.golang.org/grpc/status"
    23  
    24  	"go.chromium.org/luci/client/casclient"
    25  	"go.chromium.org/luci/common/errors"
    26  	"go.chromium.org/luci/grpc/grpcutil"
    27  	"go.chromium.org/luci/server/auth"
    28  	"go.chromium.org/luci/server/router"
    29  )
    30  
    31  // clientCacheKey is a context key type for ClientCache value.
    32  type clientCacheKey struct{}
    33  
    34  // ccKey is a context key for ClientCache value.
    35  var ccKey = &clientCacheKey{}
    36  
    37  // ClientCache caches CAS clients, one per an instance.
    38  type ClientCache struct {
    39  	lock    sync.RWMutex
    40  	clients map[string]*client.Client
    41  	ctx     context.Context
    42  }
    43  
    44  // NewClientCache initializes ClientCache.
    45  func NewClientCache(ctx context.Context) *ClientCache {
    46  	return &ClientCache{
    47  		clients: make(map[string]*client.Client),
    48  		ctx:     ctx,
    49  	}
    50  }
    51  
    52  // Get returns a Client by loading it from cache or creating a new one.
    53  func (cc *ClientCache) Get(instance string) (*client.Client, error) {
    54  	// Load Client from cache.
    55  	cc.lock.RLock()
    56  	cl, ok := cc.clients[instance]
    57  	cc.lock.RUnlock()
    58  
    59  	if ok {
    60  		return cl, nil
    61  	}
    62  
    63  	cc.lock.Lock()
    64  	defer cc.lock.Unlock()
    65  
    66  	// Somebody may have already set a client for the same instance.
    67  	cl, ok = cc.clients[instance]
    68  	if ok {
    69  		return cl, nil
    70  	}
    71  
    72  	cl, err := NewClient(cc.ctx, instance)
    73  	if err != nil {
    74  		return nil, err
    75  	}
    76  
    77  	// Cache the client.
    78  	cc.clients[instance] = cl
    79  	return cl, nil
    80  }
    81  
    82  // Clear closes Clients gracefully, and removes them from cache.
    83  func (cc *ClientCache) Clear() {
    84  	cc.lock.Lock()
    85  	defer cc.lock.Unlock()
    86  	for inst, cl := range cc.clients {
    87  		cl.Close()
    88  		delete(cc.clients, inst)
    89  	}
    90  }
    91  
    92  // withClientCacheMW creates a middleware that injects the ClientCache to context.
    93  func withClientCacheMW(cc *ClientCache) router.Middleware {
    94  	return func(c *router.Context, next router.Handler) {
    95  		c.Request = c.Request.WithContext(context.WithValue(c.Request.Context(), ccKey, cc))
    96  		next(c)
    97  	}
    98  }
    99  
   100  // GetClient returns a Client by loading it from cache or creating a new one.
   101  func GetClient(c context.Context, instance string) (*client.Client, error) {
   102  	cc, err := clientCache(c)
   103  	if err != nil {
   104  		return nil, err
   105  	}
   106  	return cc.Get(instance)
   107  }
   108  
   109  // clientCache returns ClientCache by retrieving it from the context.
   110  func clientCache(c context.Context) (*ClientCache, error) {
   111  	cc, ok := c.Value(ccKey).(*ClientCache)
   112  	if !ok {
   113  		return nil, errors.New("ClientCache not installed in the context")
   114  	}
   115  	return cc, nil
   116  }
   117  
   118  // NewClient connects to the instance of remote execution service, and returns a client.
   119  func NewClient(ctx context.Context, instance string) (*client.Client, error) {
   120  	creds, err := auth.GetPerRPCCredentials(ctx, auth.AsSelf, auth.WithScopes(auth.CloudOAuthScopes...))
   121  	if err != nil {
   122  		return nil, errors.Annotate(err, "failed to get credentials").Err()
   123  	}
   124  
   125  	c, err := client.NewClient(ctx, instance,
   126  		client.DialParams{
   127  			Service:            casclient.AddrProd,
   128  			TransportCredsOnly: true,
   129  		},
   130  		&client.PerRPCCreds{Creds: creds},
   131  		client.StartupCapabilities(false),
   132  	)
   133  	if err != nil {
   134  		// convert gRPC code to LUCI errors tag.
   135  		t := grpcutil.Tag.With(status.Code(err))
   136  		return nil, errors.Annotate(err, "failed to create client").Tag(t).Err()
   137  	}
   138  
   139  	return c, nil
   140  }