Sharing data between multiple java threads and get the updated value

Question

I want to create a java application , where we want to make rest calls for multiple users , with the help of an access token. I am using 1 thread per user . The access token, that I am using , is valid for 1 hour.Once the token expires , I will get an 401 error , and have to update the token for all the threads , and continue. I am thinking of using a volatile variable which I have made static to update all the threads. My requirement , is , the moment I get to know in one of threads that the token has expired , I want all the threads to stop processing and wait till the new token is generated(this takes couple of seconds).Also once generated , the token should be updated automatically , without each thread failing due to expired token.

Below is a sample code that I have written :

import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

public class Sample {

public static void main(String[] args) {

    String[] myStrings = { "User1" , "User2" , "User3" };

    ScheduledExecutorService scheduledExecutorService = Executors
            .newScheduledThreadPool(myStrings.length);

    TokenGenerator.getToken();

    for(String str : myStrings){
        scheduledExecutorService.scheduleAtFixedRate(new Task(str), 0, 5, TimeUnit.SECONDS);    
    }
}

}

 class Task implements Runnable{

private String name;

public Task(String name){
    this.name = name;

}


@Override
public void run() {

    getResponse(TokenGenerator.token);

}

private void getResponse(String token) {
    // Make http calls
    // if token expire , call getToken again. Pause all the running threads , and
    // update the token for all threads

    TokenGenerator.getToken();
}

 }

 class TokenGenerator {

public static volatile String token;

public static void getToken() {

    token = "new Token everytime";

}

 }

Is there a better approach to this problem? The above code does not satisfies my use case, since once a thread starts generating a new token, all other threads are not being paused. Request to suggest some improvements..

Answer 1

You can put the token in an AtomicReference and use a Semaphore to pause the threads:

public class TokenWrapper {
  private final AtomicReference<Token> tokenRef = new AtomicReference<>(null);
  private final Semaphore semaphore = new Semaphore(Integer.MAX_VALUE);

  public TokenWrapper() {
    Token newToken = // refresh token
    tokenRef.set(newToken);
  }

  public Token getToken() {
    Token token = null;
    while((token = tokenRef.get()) == null) {
      semaphore.acquire();
    }
    return token;
  }

  public Token refreshToken(Token oldToken) {
    if(tokenRef.compareAndSet(oldToken, null)) {
      semaphore.drainPermits();          
      Token newToken = // refresh token
      tokenRef.set(newToken);
      semaphore.release(Integer.MAX_VALUE);
      return newToken;
    } else return getToken();
  }
}

public class RESTService {
  private static final TokenWrapper tokenWrapper = new TokenWrapper();

  public void run() {
    Token token = tokenWrapper.getToken();
    Response response = // call service with token
    if(response.getStatus == 401) {
      tokenWrapper.refreshToken(token);
    }
  }
}

refreshToken() uses an atomic compareAndSet on tokenRef to ensure that only one thread will refresh the token, and then calls drainPermits() on the semaphore to cause other threads to wait until the token is refreshed. getToken() returns the token if it isn't null, else waits on the semaphore - this is done in a loop because it's possible that a thread will have to spin for a few cycles between tokenRef being set to null and drainPermits() being called on semaphore.

---

Edit: Modified the signature of refreshToken(Token oldToken) so that the old token is passed in rather than being read inside of the method - this is to prevent a situation where RESTService_A refreshes the token, RESTService_B gets a 401 with the old expired token, and then RESTService_B calls refreshToken after RESTService_A's call to refreshToken has completed, resulting in the token being refreshed twice. With the new signature, RESTService_B will pass in the old expired token, and so the compareAndSet call will fail when the old token fails to match the new token, resulting in refreshToken only being called once.