Building a basic token manager

27th April, 2024 (64 days ago)

#golang#concurrency

For the past few days I have been regularly working on integrating API’s and most of us developers do that on regular basis.

The most common thing in all these API’s is (not bad documentation 😅) getting an access token to access protected API’s which expires in a certain period of time.

Problem Statement I came across:

Generate an access token after some interval and store it
Should be easily accessible as it will be required again and again

Most suggested way to go about this based on discussion with my team was :

Push it to redis and add some TTL(time to live), if the token’s there use it or generate it

Being a gopher for a long time now. Whenever I think of an async task, I think about go routines

So I came up with an idea of a simple token manager.

// TokenGeneratorFunc is a function type that takes no
// arguments and returns an interface{} value
// and an error value. This function type
// is used to generate tokens.
type TokenGeneratorFunc func() (interface{}, error)

// TokenManager is a struct that manages
// the generation and caching of tokens.
type TokenManager struct {
    // token is an atomic.Value that stores
    // the current token value.
    // It uses atomic operations to
    // ensure thread-safety when accessing
    // or modifying the token value concurrently.
    token atomic.Value

    // duration specifies
    // the time duration for which a token
    // is valid and should be
    // generated duration time.Duration
    duration time.Duration


    // generatorFunc is a function of
    // type TokenGeneratorFunc
    // that generates a new token
    // when called.
    generatorFunc TokenGeneratorFunc

    // name is a string that
    // identifies the
    // TokenManager instance.
    name string

    // hasGeneratedFirstToken is a bool that
    // indicates whether the first token
    // has been generated or not.
    // It is used to handle
    // the first token generation
    // differently, if required.
    hasGeneratedFirstToken bool

    // blockFirstTime is a channel that
    // can be used to block the first
    // token generation until a signal
    // is received on the channel.
    // This can be useful in scenarios
    // where the first token generation
    // needs to be delayed or
    // synchronized with other operations.
    blockFirstTime chan bool

}

The TokenManager has a few key responsibilities:

Token Generation: It takes a TokenGeneratorFunc as an input, which is a function that generates a new token. This function could be making an API call, performing some cryptographic operation, or any other logic required to generate a token.
Token Refresh: The TokenManager periodically calls the TokenGeneratorFunc to refresh the token. The refresh interval is determined by the duration field, which specifies the time between token refreshes.
Token Storage: The TokenManager stores the current token in an atomic variable, ensuring thread-safety when multiple goroutines try to access or update the token concurrently.
Token Distribution: Other parts of the system can retrieve the current token by calling the GetToken method on the TokenManager.

Let’s get into the code now

Constructor - instantiates the token manager

func NewTokenManager(
    name string,
    duration *time.Duration,
    generatorFunc TokenGeneratorFunc) *TokenManager {
    tm := &TokenManager{
        duration: *duration,
        generatorFunc: generatorFunc,
        name: name,
        // making the channel size 2 so it
        // doesn't block the first write
        blockFirstTime: make(chan bool, 2),
        hasGeneratedFirstToken: false,
    }
    tm.token.Store("")
    return tm
}

This is a constructor function that creates a new TokenManager instance. It takes three arguments:

name: The name of the resource or service for which the token is being managed.
duration: A pointer to a time.Duration that specifies the interval between token refreshes.
generatorFunc: The TokenGeneratorFunc that will be used to generate new tokens.

The function initializes the TokenManager struct with the provided values, creates a new channel blockFirstTime, and stores an empty string as the initial token value.

GetToken - retrieves token

func (t *TokenManager) GetToken() interface{} {
   // If the first token hasn't been generated yet
   if !t.hasGeneratedFirstToken {
       // Block until a value is received from
       // the blockFirstTime channel
       <-t.blockFirstTime
   }

   // Return the current token value stored
   // in the atomic.Value
   return t.token.Load()
}

The GetToken method is used to retrieve the current token from the TokenManager,.

If the first token hasn’t been generated yet, it blocks until the blockFirstTime, channel receives a value (which happens after the first token is generated).

Then, it returns the current token by loading it from the token field using the atomic.Value.Load, method.

RunTokenGenerator - runs an async task to generate tokens

func (tm *TokenManager) RunTokenGenerator() {
   // Create a new ticker that fires at
   // intervals specified by tm.duration
   ticker := time.NewTicker(tm.duration)
   // Ensure the ticker is stopped
   // when this function returns
   defer ticker.Stop()

   // Generate and update the initial token
   tm.UpdateToken()

   // This loop will run indefinitely
   for range ticker.C {
       // On each tick, generate and update the token
       tm.UpdateToken()
   }
}

The RunTokenGenerator method is responsible for periodically refreshing the token.

It creates a new time.Ticker based on the duration field, which sends a value on the ticker.C channel at the specified interval.

The method immediately generates the first token by calling tm.UpdateToken(). Then, it enters a loop that blocks until a value is received on the ticker.C channel, at which point it calls tm.UpdateToken() again to refresh the token.

func (tm *TokenManager) UpdateToken() {
   // Print a message indicating that a new
   // token is being generated
   fmt.Println("Generating new session token
   for connecting to " + tm.name)

   // Call the token generator function
   // to generate a new token
   response, err := tm.generatorFunc()
   if err != nil {
       // If there was an error generating
       // the token, print the error and return
       fmt.Println("Error updating token:", err)
       return
   }

   // Store the newly generated token
   //in the atomic.Value
   tm.token.Store(response)

   // If this is the first time a token is being generated
   if !tm.hasGeneratedFirstToken {
       // Set the hasGeneratedFirstToken flag to true
       tm.hasGeneratedFirstToken = true

       // Send a value to the blockFirstTime channel
       // to unblock any goroutines waiting for
       // the first token generation
       tm.blockFirstTime <- true
   }

   // Print a message indicating that the
   // new token was successfully generated
   fmt.Println("Successfully generated new
    token for " + tm.name)
}

The UpdateToken method is responsible for generating a new token and storing it in the TokenManager. It first prints a log message indicating that it’s generating a new token for the specified name.

Then, it calls the generatorFunc to generate a new token. If an error occurs during token generation, it prints the error and returns without updating the token. If the token is generated successfully, it stores the new token in the token field using the atomic.Value.Store method.

If this is the first token being generated, it sets the hasGeneratedFirstToken flag to true and sends a value on the blockFirstTime channel, allowing the GetToken method to unblock and return the newly generated token.

Checkout the full code at Github

LinkedIn Github Twitter