github.com/myhau/pulumi/pkg/v3@v3.70.2-0.20221116134521-f2775972e587/backend/httpstate/token_source.go (about)

     1  // Copyright 2016-2022, Pulumi Corporation.
     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 httpstate
    16  
    17  import (
    18  	"context"
    19  	"fmt"
    20  	"time"
    21  
    22  	"github.com/pulumi/pulumi/sdk/v3/go/common/util/logging"
    23  )
    24  
    25  type tokenSourceCapability interface {
    26  	GetToken() (string, error)
    27  }
    28  
    29  // tokenSource is a helper type that manages the renewal of the lease token for a managed update.
    30  type tokenSource struct {
    31  	requests chan tokenRequest
    32  	done     chan bool
    33  }
    34  
    35  var _ tokenSourceCapability = &tokenSource{}
    36  
    37  type tokenRequest chan<- tokenResponse
    38  
    39  type tokenResponse struct {
    40  	token string
    41  	err   error
    42  }
    43  
    44  func newTokenSource(
    45  	ctx context.Context,
    46  	initialToken string,
    47  	initialTokenExpires time.Time,
    48  	duration time.Duration,
    49  	refreshToken func(
    50  		ctx context.Context,
    51  		duration time.Duration,
    52  		currentToken string,
    53  	) (string, time.Time, error),
    54  ) (*tokenSource, error) {
    55  	requests, done := make(chan tokenRequest), make(chan bool)
    56  	ts := &tokenSource{requests: requests, done: done}
    57  	go ts.handleRequests(ctx, initialToken, initialTokenExpires, duration, refreshToken)
    58  	return ts, nil
    59  }
    60  
    61  func (ts *tokenSource) handleRequests(
    62  	ctx context.Context,
    63  	initialToken string,
    64  	initialTokenExpires time.Time,
    65  	duration time.Duration,
    66  	refreshToken func(
    67  		ctx context.Context,
    68  		duration time.Duration,
    69  		currentToken string,
    70  	) (string, time.Time, error),
    71  ) {
    72  
    73  	renewTicker := time.NewTicker(duration / 8)
    74  	defer renewTicker.Stop()
    75  
    76  	state := struct {
    77  		token   string    // most recently renewed token
    78  		error   error     // non-nil indicates a terminal error state
    79  		expires time.Time // assumed expiry of the token
    80  	}{
    81  		token:   initialToken,
    82  		expires: initialTokenExpires,
    83  	}
    84  
    85  	renewUpdateLeaseIfStale := func() {
    86  		if state.error != nil {
    87  			return
    88  		}
    89  
    90  		now := time.Now()
    91  
    92  		// We will renew the lease after 50% of the duration
    93  		// has elapsed to allow time for retries.
    94  		stale := now.Add(duration / 2).After(state.expires)
    95  		if !stale {
    96  			return
    97  		}
    98  
    99  		newToken, newTokenExpires, err := refreshToken(ctx, duration, state.token)
   100  		// If renew failed, all further GetToken requests will return this error.
   101  		if err != nil {
   102  			logging.V(3).Infof("error renewing lease: %v", err)
   103  			state.error = fmt.Errorf("renewing lease: %w", err)
   104  			renewTicker.Stop()
   105  		} else {
   106  			state.token = newToken
   107  			state.expires = newTokenExpires
   108  		}
   109  	}
   110  
   111  	for {
   112  		select {
   113  		case <-renewTicker.C:
   114  			renewUpdateLeaseIfStale()
   115  		case c, ok := <-ts.requests:
   116  			if !ok {
   117  				close(ts.done)
   118  				return
   119  			}
   120  			// If ticker has not kept up, block on
   121  			// renewing rather than risking returning a
   122  			// stale token.
   123  			renewUpdateLeaseIfStale()
   124  			if state.error == nil {
   125  				c <- tokenResponse{token: state.token}
   126  			} else {
   127  				c <- tokenResponse{err: state.error}
   128  			}
   129  		}
   130  	}
   131  }
   132  
   133  func (ts *tokenSource) GetToken() (string, error) {
   134  	ch := make(chan tokenResponse)
   135  	ts.requests <- ch
   136  	resp := <-ch
   137  	return resp.token, resp.err
   138  }
   139  
   140  func (ts *tokenSource) Close() {
   141  	close(ts.requests)
   142  	<-ts.done
   143  }