vitess.io/vitess@v0.16.2/go/vt/vtadmin/cache/refresh.go (about)

     1  /*
     2  Copyright 2022 The Vitess Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package cache
    18  
    19  import (
    20  	"context"
    21  	"fmt"
    22  	"net/http"
    23  	"strconv"
    24  	"strings"
    25  
    26  	"google.golang.org/grpc/metadata"
    27  
    28  	"vitess.io/vitess/go/vt/log"
    29  )
    30  
    31  var (
    32  	cacheRefreshHeader          string
    33  	cacheRefreshGRPCMetadataKey string
    34  )
    35  
    36  // SetCacheRefreshKey sets the global HTTP header and gRPC metadata keys for
    37  // requests to force a cache refresh. Any whitespace characters are replaced
    38  // with hyphens.
    39  //
    40  // It is not threadsafe, and should be called only at startup or in an init
    41  // function.
    42  func SetCacheRefreshKey(k string) {
    43  	l := strings.ToLower(k)
    44  	l = strings.ReplaceAll(l, " ", "-")
    45  	cacheRefreshHeader, cacheRefreshGRPCMetadataKey = fmt.Sprintf("x-%s", l), l
    46  }
    47  
    48  // NewIncomingRefreshContext returns an incoming gRPC context with metadata
    49  // set to signal a cache refresh.
    50  func NewIncomingRefreshContext(ctx context.Context) context.Context {
    51  	md := metadata.Pairs(cacheRefreshGRPCMetadataKey, "true")
    52  	return metadata.NewIncomingContext(ctx, md)
    53  }
    54  
    55  // ShouldRefreshFromIncomingContext returns true if the gRPC metadata in the
    56  // incoming context signals a cache refresh was requested.
    57  func ShouldRefreshFromIncomingContext(ctx context.Context) bool {
    58  	if cacheRefreshGRPCMetadataKey == "" {
    59  		return false
    60  	}
    61  
    62  	md, ok := metadata.FromIncomingContext(ctx)
    63  	if !ok {
    64  		return false
    65  	}
    66  
    67  	vals := md.Get(cacheRefreshGRPCMetadataKey)
    68  	if len(vals) == 0 {
    69  		return false
    70  	}
    71  
    72  	shouldRefresh, err := strconv.ParseBool(vals[0])
    73  	if err != nil {
    74  		log.Warningf("failed to parse %s metadata key as bool: %s", cacheRefreshGRPCMetadataKey, err)
    75  		return false
    76  	}
    77  
    78  	return shouldRefresh
    79  }
    80  
    81  // ShouldRefreshFromRequest returns true if the HTTP request headers signal a
    82  // cache refresh was requested.
    83  func ShouldRefreshFromRequest(r *http.Request) bool {
    84  	if cacheRefreshHeader == "" {
    85  		return false
    86  	}
    87  
    88  	h := r.Header.Get(cacheRefreshHeader)
    89  	if h == "" {
    90  		return false
    91  	}
    92  
    93  	shouldRefresh, err := strconv.ParseBool(h)
    94  	if err != nil {
    95  		log.Warningf("failed to parse %s header as bool: %s", cacheRefreshHeader, err)
    96  		return false
    97  	}
    98  
    99  	return shouldRefresh
   100  }