discussion: token refresh mechanism for rest client

Background: #301

The token fetched from the token server may have a TTL, see TokenResponse::expires_in. In most implementations, the value is about several hours. Our catalog client is a long-lived object, which means that we should handle the token expiration event.

There are two ways:

Spawn a background task, and setup a ticker with a specified time interval, e.g. expires_in * 0.9 seconds, and refetch the token when triggered.

Pros: Easy to implement
Cons: Must trust the local timer skew
Cons: Must introduce a timer, which means depending on a specified async runtime

Call every methods with a retry wrapper. When meeting an unauthorized error, refetch the token and retry the method.

Pros: Consistent with iceberg-python
Pros: Does not rely on local clock and specific runtime
Cons: When expired, thousands of concurrent requests may fail, and then all of them will trigger a token refetch, which is not ideal.
- This can be workaround by some concurrency control, to force only one request will refetch the token and others must wait for the result, but it introduced complexity.

2. Call every methods with a retry wrapper. When meeting an unauthorized error, refetch the token and retry the method.

Our rest catalog client now has a authenticate function.

iceberg-rust/crates/catalog/rest/src/client.rs

Line 88 in 0937f19

async fn authenticate(&self, req: &mut Request) -> Result<()> {

We can perform the expiration logic here.

I also prefer 2. The iceberg community may deprecate current auth approach in future, and introduce true oauth client in future. I prefer to have a simpler solution without introducing too much extra dependency.

Store (expired_at, token) pair in an async Mutex, and check whether it's close to expiration when needed. If it's close to expiration, lock the mutex, update the pair, and unlock. Then only one request will update the token.

3. Store (expired_at, token) pair in an async Mutex, and check whether it's close to expiration when needed. If it's close to expiration, lock the mutex, update the pair, and unlock. Then only one request will update the token.

Great!

Store (expired_at, token) pair in an async Mutex, and check whether it's close to expiration when needed. If it's close to expiration, lock the mutex, update the pair, and unlock. Then only one request will update the token.

Sounds reasonable to me.

Then we need https://docs.rs/tokio/latest/tokio/sync/struct.RwLock.html, still depends on tokio :(

This is rust

Then we need https://docs.rs/tokio/latest/tokio/sync/struct.RwLock.html, still depends on tokio :(

This is rust

How about futures Mutex?

The current implementation utilizes std::sync::Mutex with a precisely controlled locking scope.

iceberg-rust/crates/catalog/rest/src/client.rs

Lines 88 to 179 in 8562708

    
               async fn authenticate(&self, req: &mut Request) -> Result<()> { 
        
                   // Clone the token from lock without holding the lock for entire function. 
        
                   let token = { self.token.lock().expect("lock poison").clone() }; 
        
                   if self.credential.is_none() && token.is_none() { 
        
                       return Ok(()); 
        
                   } 
        
                   // Use token if provided. 
        
                   if let Some(token) = &token { 
        
                       req.headers_mut().insert( 
        
                           http::header::AUTHORIZATION, 
        
                           format!("Bearer {token}").parse().map_err(|e| { 
        
                               Error::new( 
        
                                   ErrorKind::DataInvalid, 
        
                                   "Invalid token received from catalog server!", 
        
                               ) 
        
                               .with_source(e) 
        
                           })?, 
        
                       ); 
        
                       return Ok(()); 
        
                   } 
        
                   // Credential must exist here. 
        
                   let (client_id, client_secret) = self.credential.as_ref().ok_or_else(|| { 
        
                       Error::new( 
        
                           ErrorKind::DataInvalid, 
        
                           "Credential must be provided for authentication", 
        
                       ) 
        
                   })?; 
        
                   let mut params = HashMap::with_capacity(4); 
        
                   params.insert("grant_type", "client_credentials"); 
        
                   if let Some(client_id) = client_id { 
        
                       params.insert("client_id", client_id); 
        
                   } 
        
                   params.insert("client_secret", client_secret); 
        
                   params.extend( 
        
                       self.extra_oauth_params 
        
                           .iter() 
        
                           .map(|(k, v)| (k.as_str(), v.as_str())), 
        
                   ); 
        
                   let auth_req = self 
        
                       .client 
        
                       .request(Method::POST, &self.token_endpoint) 
        
                       .form(&params) 
        
                       .build()?; 
        
                   let auth_resp = self.client.execute(auth_req).await?; 
        
                   let auth_res: TokenResponse = if auth_resp.status().as_u16() == OK { 
        
                       let text = auth_resp.bytes().await?; 
        
                       Ok(serde_json::from_slice(&text).map_err(|e| { 
        
                           Error::new( 
        
                               ErrorKind::Unexpected, 
        
                               "Failed to parse response from rest catalog server!", 
        
                           ) 
        
                           .with_context("json", String::from_utf8_lossy(&text)) 
        
                           .with_source(e) 
        
                       })?) 
        
                   } else { 
        
                       let code = auth_resp.status(); 
        
                       let text = auth_resp.bytes().await?; 
        
                       let e: ErrorResponse = serde_json::from_slice(&text).map_err(|e| { 
        
                           Error::new( 
        
                               ErrorKind::Unexpected, 
        
                               "Failed to parse response from rest catalog server!", 
        
                           ) 
        
                           .with_context("json", String::from_utf8_lossy(&text)) 
        
                           .with_context("code", code.to_string()) 
        
                           .with_source(e) 
        
                       })?; 
        
                       Err(Error::from(e)) 
        
                   }?; 
        
                   let token = auth_res.access_token; 
        
                   // Update token. 
        
                   *self.token.lock().expect("lock poison") = Some(token.clone()); 
        
                   // Insert token in request. 
        
                   req.headers_mut().insert( 
        
                       http::header::AUTHORIZATION, 
        
                       format!("Bearer {token}").parse().map_err(|e| { 
        
                           Error::new( 
        
                               ErrorKind::DataInvalid, 
        
                               "Invalid token received from catalog server!", 
        
                           ) 
        
                           .with_source(e) 
        
                       })?, 
        
                   ); 
        
                   Ok(()) 
        
               }

The worst thing that could happen here is that we might send multiple token refresh requests, which I think is fine for a catalog.

It's acceptable, but not as expected. When a client request a token, the server may expect it will be used for several hours.

as an alternative, can we simply expose a public refresh_token method that allows the caller to implement their own token invalidation (for my use case, via option 1, using tokio)

this avoids the downside of "thousands of concurrent requests may fail", and does not require the core library to take any additional dependencies

this change perhaps does not remove the need to implement option 2 in this library, but it does provide an option for downstream consumers to implement token expiry as they wish (where there currently is no such option)

@liurenjie1024

@liurenjie1024 I've implemented this here: #1465

feat(iceberg-catalog-rest): expose invalidate_token, regenerate_token…

@liurenjie1024 I've implemented this here: #1465

Thanks, we could close this for now.

Navigation Menu

discussion: token refresh mechanism for rest client #437

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

	async fn authenticate(&self, req: &mut Request) -> Result<()> {
	// Clone the token from lock without holding the lock for entire function.
	let token = { self.token.lock().expect("lock poison").clone() };

	if self.credential.is_none() && token.is_none() {
	return Ok(());
	}

	// Use token if provided.
	if let Some(token) = &token {
	req.headers_mut().insert(
	http::header::AUTHORIZATION,
	format!("Bearer {token}").parse().map_err(\|e\| {
	Error::new(
	ErrorKind::DataInvalid,
	"Invalid token received from catalog server!",
	)
	.with_source(e)
	})?,
	);
	return Ok(());
	}

	// Credential must exist here.
	let (client_id, client_secret) = self.credential.as_ref().ok_or_else(\|\| {
	Error::new(
	ErrorKind::DataInvalid,
	"Credential must be provided for authentication",
	)
	})?;

	let mut params = HashMap::with_capacity(4);
	params.insert("grant_type", "client_credentials");
	if let Some(client_id) = client_id {
	params.insert("client_id", client_id);
	}
	params.insert("client_secret", client_secret);
	params.extend(
	self.extra_oauth_params
	.iter()
	.map(\|(k, v)\| (k.as_str(), v.as_str())),
	);

	let auth_req = self
	.client
	.request(Method::POST, &self.token_endpoint)
	.form(&params)
	.build()?;
	let auth_resp = self.client.execute(auth_req).await?;

	let auth_res: TokenResponse = if auth_resp.status().as_u16() == OK {
	let text = auth_resp.bytes().await?;
	Ok(serde_json::from_slice(&text).map_err(\|e\| {
	Error::new(
	ErrorKind::Unexpected,
	"Failed to parse response from rest catalog server!",
	)
	.with_context("json", String::from_utf8_lossy(&text))
	.with_source(e)
	})?)
	} else {
	let code = auth_resp.status();
	let text = auth_resp.bytes().await?;
	let e: ErrorResponse = serde_json::from_slice(&text).map_err(\|e\| {
	Error::new(
	ErrorKind::Unexpected,
	"Failed to parse response from rest catalog server!",
	)
	.with_context("json", String::from_utf8_lossy(&text))
	.with_context("code", code.to_string())
	.with_source(e)
	})?;
	Err(Error::from(e))
	}?;
	let token = auth_res.access_token;
	// Update token.
	*self.token.lock().expect("lock poison") = Some(token.clone());
	// Insert token in request.
	req.headers_mut().insert(
	http::header::AUTHORIZATION,
	format!("Bearer {token}").parse().map_err(\|e\| {
	Error::new(
	ErrorKind::DataInvalid,
	"Invalid token received from catalog server!",
	)
	.with_source(e)
	})?,
	);

	Ok(())
	}

discussion: token refresh mechanism for rest client #437

Description

Activity

Xuanwo commented on Jul 5, 2024

liurenjie1024 commented on Jul 5, 2024

TennyZhuang commented on Jul 5, 2024

Xuanwo commented on Jul 5, 2024

liurenjie1024 commented on Jul 5, 2024

TennyZhuang commented on Jul 5, 2024

liurenjie1024 commented on Jul 5, 2024

Xuanwo commented on Jul 5, 2024

TennyZhuang commented on Jul 5, 2024

cmcarthur commented on Jun 23, 2025

cmcarthur commented on Jun 24, 2025

liurenjie1024 commented on Jun 25, 2025

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions