Skip to content

Commit f0ca0f6

Browse files
Improve Tower client caching via Tiered cache (#772)
Signed-off-by: Paolo Di Tommaso <[email protected]> Signed-off-by: munishchouhan <[email protected]> Co-authored-by: munishchouhan <[email protected]>
1 parent 13a5937 commit f0ca0f6

32 files changed

+1101
-118
lines changed

src/main/groovy/io/seqera/wave/encoder/MoshiEncodeStrategy.groovy

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -54,13 +54,18 @@ abstract class MoshiEncodeStrategy<V> implements EncodingStrategy<V> {
5454
init()
5555
}
5656

57+
MoshiEncodeStrategy(JsonAdapter.Factory customFactory) {
58+
this.type = TypeHelper.getGenericType(this, 0)
59+
init(customFactory)
60+
}
61+
5762
MoshiEncodeStrategy(Type type) {
5863
this.type = type
5964
init()
6065
}
6166

62-
private void init() {
63-
this.moshi = new Moshi.Builder()
67+
private void init(JsonAdapter.Factory customFactory=null) {
68+
final builder = new Moshi.Builder()
6469
.add(new ByteArrayAdapter())
6570
.add(new DateTimeAdapter())
6671
.add(new PathAdapter())
@@ -73,9 +78,11 @@ abstract class MoshiEncodeStrategy<V> implements EncodingStrategy<V> {
7378
.withSubtype(ProxyHttpRequest.class, ProxyHttpRequest.simpleName)
7479
.withSubtype(ProxyHttpResponse.class, ProxyHttpResponse.simpleName)
7580
.withSubtype(PairingHeartbeat.class, PairingHeartbeat.simpleName)
76-
.withSubtype(PairingResponse.class, PairingResponse.simpleName)
77-
)
78-
.build()
81+
.withSubtype(PairingResponse.class, PairingResponse.simpleName) )
82+
// add custom factory if provider
83+
if( customFactory )
84+
builder.add(customFactory)
85+
this.moshi = builder.build()
7986
this.jsonAdapter = moshi.adapter(type)
8087

8188
}
Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
/*
2+
* Wave, containers provisioning service
3+
* Copyright (c) 2023-2024, Seqera Labs
4+
*
5+
* This program is free software: you can redistribute it and/or modify
6+
* it under the terms of the GNU Affero General Public License as published by
7+
* the Free Software Foundation, either version 3 of the License, or
8+
* (at your option) any later version.
9+
*
10+
* This program is distributed in the hope that it will be useful,
11+
* but WITHOUT ANY WARRANTY; without even the implied warranty of
12+
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
13+
* GNU Affero General Public License for more details.
14+
*
15+
* You should have received a copy of the GNU Affero General Public License
16+
* along with this program. If not, see <https://www.gnu.org/licenses/>.
17+
*/
18+
19+
package io.seqera.wave.encoder
20+
21+
/**
22+
* Marker interface for Moshi encoded exchange objects
23+
*
24+
* @author Paolo Di Tommaso <[email protected]>
25+
*/
26+
interface MoshiExchange {
27+
}

src/main/groovy/io/seqera/wave/service/CredentialServiceImpl.groovy

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ class CredentialServiceImpl implements CredentialsService {
6060
if (pairing.isExpired())
6161
log.debug("Exchange key registered for service ${PairingService.TOWER_SERVICE} at endpoint: ${identity.towerEndpoint} used after expiration, should be renewed soon")
6262

63-
final all = towerClient.listCredentials(identity.towerEndpoint, JwtAuth.of(identity), identity.workspaceId).get().credentials
63+
final all = towerClient.listCredentials(identity.towerEndpoint, JwtAuth.of(identity), identity.workspaceId, identity.workflowId).credentials
6464

6565
if (!all) {
6666
log.debug "No credentials found for userId=$identity.userId; workspaceId=$identity.workspaceId; endpoint=$identity.towerEndpoint"
@@ -92,7 +92,7 @@ class CredentialServiceImpl implements CredentialsService {
9292
// log for debugging purposes
9393
log.debug "Credentials matching criteria registryName=$registryName; userId=$identity.userId; workspaceId=$identity.workspaceId; endpoint=$identity.towerEndpoint => $creds"
9494
// now fetch the encrypted key
95-
final encryptedCredentials = towerClient.fetchEncryptedCredentials(identity.towerEndpoint, JwtAuth.of(identity), creds.id, pairing.pairingId, identity.workspaceId).get()
95+
final encryptedCredentials = towerClient.fetchEncryptedCredentials(identity.towerEndpoint, JwtAuth.of(identity), creds.id, pairing.pairingId, identity.workspaceId, identity.workflowId)
9696
final privateKey = pairing.privateKey
9797
final credentials = decryptCredentials(privateKey, encryptedCredentials.keys)
9898
return parsePayload(credentials)
@@ -112,7 +112,7 @@ class CredentialServiceImpl implements CredentialsService {
112112
final response = towerClient.describeWorkflowLaunch(identity.towerEndpoint, JwtAuth.of(identity), identity.workflowId)
113113
if( !response )
114114
return null
115-
final computeEnv = response.get()?.launch?.computeEnv
115+
final computeEnv = response?.launch?.computeEnv
116116
if( !computeEnv )
117117
return null
118118
if( computeEnv.platform != 'aws-batch' )

src/main/groovy/io/seqera/wave/service/UserServiceImpl.groovy

Lines changed: 14 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -19,17 +19,18 @@
1919
package io.seqera.wave.service
2020

2121
import java.util.concurrent.CompletableFuture
22-
import java.util.concurrent.ExecutionException
23-
import io.micronaut.core.annotation.Nullable
22+
import java.util.concurrent.ExecutorService
2423

2524
import groovy.transform.CompileStatic
2625
import groovy.util.logging.Slf4j
26+
import io.micronaut.core.annotation.Nullable
27+
import io.micronaut.scheduling.TaskExecutors
2728
import io.seqera.wave.exception.UnauthorizedException
2829
import io.seqera.wave.tower.User
2930
import io.seqera.wave.tower.auth.JwtAuth
3031
import io.seqera.wave.tower.client.TowerClient
31-
import io.seqera.wave.tower.client.UserInfoResponse
3232
import jakarta.inject.Inject
33+
import jakarta.inject.Named
3334
import jakarta.inject.Singleton
3435
/**
3536
* Define a service to access a Tower user
@@ -45,28 +46,24 @@ class UserServiceImpl implements UserService {
4546
@Nullable
4647
private TowerClient towerClient
4748

49+
@Inject
50+
@Named(TaskExecutors.BLOCKING)
51+
private ExecutorService ioExecutor
52+
4853
@Override
4954
CompletableFuture<User> getUserByAccessTokenAsync(String endpoint, JwtAuth auth) {
5055
if( !towerClient )
5156
throw new IllegalStateException("Missing Tower client - make sure the 'tower' micronaut environment has been provided")
5257

53-
towerClient.userInfo(endpoint, auth).handle( (UserInfoResponse resp, Throwable error) -> {
54-
if( error )
55-
throw error
56-
if (!resp || !resp.user)
57-
throw new UnauthorizedException("Unauthorized - Make sure you have provided a valid access token")
58-
log.debug("Authorized user=$resp.user")
59-
return resp.user
60-
})
58+
return CompletableFuture.supplyAsync(()-> getUserByAccessToken(endpoint,auth), ioExecutor)
6159
}
6260

6361
@Override
6462
User getUserByAccessToken(String endpoint, JwtAuth auth) {
65-
try {
66-
return getUserByAccessTokenAsync(endpoint, auth).get()
67-
}
68-
catch(ExecutionException e){
69-
throw e.cause
70-
}
63+
final resp = towerClient.userInfo(endpoint, auth)
64+
if (!resp || !resp.user)
65+
throw new UnauthorizedException("Unauthorized - Make sure you have provided a valid access token")
66+
log.debug("Authorized user=$resp.user")
67+
return resp.user
7168
}
7269
}

src/main/groovy/io/seqera/wave/service/pairing/socket/PairingWebSocket.groovy

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,7 @@ class PairingWebSocket {
7979
// Register the client and the sender callback that it's needed to deliver
8080
// the message to the remote client
8181
channel.registerClient(service, endpoint, session.id,(pairingMessage) -> {
82-
log.trace "Sendind message=${pairingMessage} - endpoint: ${endpoint} [sessionId: $session.id]"
82+
log.trace "Sending message=${pairingMessage} - endpoint: ${endpoint} [sessionId: $session.id]"
8383
session .sendAsync(pairingMessage)
8484
})
8585

Lines changed: 180 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,180 @@
1+
/*
2+
* Wave, containers provisioning service
3+
* Copyright (c) 2023-2024, Seqera Labs
4+
*
5+
* This program is free software: you can redistribute it and/or modify
6+
* it under the terms of the GNU Affero General Public License as published by
7+
* the Free Software Foundation, either version 3 of the License, or
8+
* (at your option) any later version.
9+
*
10+
* This program is distributed in the hope that it will be useful,
11+
* but WITHOUT ANY WARRANTY; without even the implied warranty of
12+
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
13+
* GNU Affero General Public License for more details.
14+
*
15+
* You should have received a copy of the GNU Affero General Public License
16+
* along with this program. If not, see <https://www.gnu.org/licenses/>.
17+
*/
18+
19+
package io.seqera.wave.store.cache
20+
21+
import java.time.Duration
22+
import java.util.concurrent.TimeUnit
23+
import java.util.concurrent.locks.Lock
24+
import java.util.concurrent.locks.ReentrantLock
25+
import java.util.function.Function
26+
27+
import com.github.benmanes.caffeine.cache.AsyncCache
28+
import com.github.benmanes.caffeine.cache.Caffeine
29+
import com.github.benmanes.caffeine.cache.RemovalCause
30+
import com.github.benmanes.caffeine.cache.RemovalListener
31+
import groovy.transform.Canonical
32+
import groovy.transform.CompileStatic
33+
import groovy.util.logging.Slf4j
34+
import io.seqera.wave.encoder.EncodingStrategy
35+
import io.seqera.wave.encoder.MoshiEncodeStrategy
36+
import io.seqera.wave.encoder.MoshiExchange
37+
import org.jetbrains.annotations.Nullable
38+
/**
39+
* Implement a tiered-cache mechanism using a local caffeine cache as 1st level access
40+
* and a 2nd-level cache backed on Redis.
41+
*
42+
* This allow the use in distributed deployment. Note however strong consistently is not guaranteed.
43+
*
44+
* @author Paolo Di Tommaso <[email protected]>
45+
*/
46+
@Slf4j
47+
@CompileStatic
48+
abstract class AbstractTieredCache<V extends MoshiExchange> implements TieredCache<String,V> {
49+
50+
@Canonical
51+
static class Entry implements MoshiExchange {
52+
MoshiExchange value
53+
long expiresAt
54+
}
55+
56+
private EncodingStrategy<Entry> encoder
57+
58+
// FIXME https://github.com/seqeralabs/wave/issues/747
59+
private AsyncCache<String,V> l1
60+
61+
private final Duration ttl
62+
63+
private L2TieredCache<String,String> l2
64+
65+
private final Lock sync = new ReentrantLock()
66+
67+
AbstractTieredCache(L2TieredCache<String,String> l2, MoshiEncodeStrategy encoder, Duration duration, long maxSize) {
68+
log.info "Cache '${getName()}' config - prefix=${getPrefix()}; ttl=${duration}; max-size: ${maxSize}; l2=${l2}"
69+
this.l2 = l2
70+
this.ttl = duration
71+
this.encoder = encoder
72+
this.l1 = Caffeine.newBuilder()
73+
.expireAfterWrite(duration.toMillis(), TimeUnit.MILLISECONDS)
74+
.maximumSize(maxSize)
75+
.removalListener(removalListener0())
76+
.buildAsync()
77+
}
78+
79+
abstract protected getName()
80+
81+
abstract protected String getPrefix()
82+
83+
private RemovalListener removalListener0() {
84+
new RemovalListener() {
85+
@Override
86+
void onRemoval(@Nullable key, @Nullable value, RemovalCause cause) {
87+
log.trace "Cache '${name}' removing key=$key; value=$value; cause=$cause"
88+
}
89+
}
90+
}
91+
92+
@Override
93+
V get(String key) {
94+
getOrCompute(key, null)
95+
}
96+
97+
V getOrCompute(String key, Function<String,V> loader) {
98+
log.trace "Cache '${name}' checking key=$key"
99+
// Try L1 cache first
100+
V value = l1.synchronous().getIfPresent(key)
101+
if (value != null) {
102+
log.trace "Cache '${name}' L1 hit (a) - key=$key => value=$value"
103+
return value
104+
}
105+
106+
sync.lock()
107+
try {
108+
value = l1.synchronous().getIfPresent(key)
109+
if (value != null) {
110+
log.trace "Cache '${name}' L1 hit (b) - key=$key => value=$value"
111+
return value
112+
}
113+
114+
// Fallback to L2 cache
115+
value = l2Get(key)
116+
if (value != null) {
117+
log.trace "Cache '${name}' L2 hit - key=$key => value=$value"
118+
// Rehydrate L1 cache
119+
l1.synchronous().put(key, value)
120+
return value
121+
}
122+
123+
// still not value found, use loader function to fetch the value
124+
if( value==null && loader!=null ) {
125+
log.trace "Cache '${name}' invoking loader - key=$key"
126+
value = loader.apply(key)
127+
if( value!=null ) {
128+
l1.synchronous().put(key,value)
129+
l2Put(key,value)
130+
}
131+
}
132+
133+
log.trace "Cache '${name}' missing value - key=$key => value=${value}"
134+
// finally return the value
135+
return value
136+
}
137+
finally {
138+
sync.unlock()
139+
}
140+
}
141+
142+
@Override
143+
void put(String key, V value) {
144+
assert key!=null, "Cache key argument cannot be null"
145+
assert value!=null, "Cache value argument cannot be null"
146+
log.trace "Cache '${name}' putting - key=$key; value=${value}"
147+
l1.synchronous().put(key, value)
148+
l2Put(key, value)
149+
}
150+
151+
protected String key0(String k) { return getPrefix() + ':' + k }
152+
153+
protected V l2Get(String key) {
154+
if( l2 == null )
155+
return null
156+
157+
final raw = l2.get(key0(key))
158+
if( raw == null )
159+
return null
160+
161+
final Entry payload = encoder.decode(raw)
162+
if( System.currentTimeMillis() > payload.expiresAt ) {
163+
log.trace "Cache '${name}' L2 exipired - key=$key => value=${payload.value}"
164+
return null
165+
}
166+
return (V) payload.value
167+
}
168+
169+
protected void l2Put(String key, V value) {
170+
if( l2 != null ) {
171+
final raw = encoder.encode(new Entry(value, ttl.toMillis() + System.currentTimeMillis()))
172+
l2.put(key0(key), raw, ttl)
173+
}
174+
}
175+
176+
void invalidateAll() {
177+
l1.synchronous().invalidateAll()
178+
}
179+
180+
}
Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
/*
2+
* Wave, containers provisioning service
3+
* Copyright (c) 2023-2024, Seqera Labs
4+
*
5+
* This program is free software: you can redistribute it and/or modify
6+
* it under the terms of the GNU Affero General Public License as published by
7+
* the Free Software Foundation, either version 3 of the License, or
8+
* (at your option) any later version.
9+
*
10+
* This program is distributed in the hope that it will be useful,
11+
* but WITHOUT ANY WARRANTY; without even the implied warranty of
12+
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
13+
* GNU Affero General Public License for more details.
14+
*
15+
* You should have received a copy of the GNU Affero General Public License
16+
* along with this program. If not, see <https://www.gnu.org/licenses/>.
17+
*/
18+
19+
package io.seqera.wave.store.cache
20+
21+
import java.time.Duration
22+
23+
/**
24+
* Define the interface for 2nd level tired cache
25+
*
26+
* @author Paolo Di Tommaso <[email protected]>
27+
*/
28+
interface L2TieredCache<K,V> extends TieredCache<K,V> {
29+
30+
31+
/**
32+
* Add a value in the cache with the specified key. If a value already exists is overridden
33+
* with the new value.
34+
*
35+
* @param key The key of the value to be added. {@code null} is not allowed.
36+
* @param value The value to be added in the cache for the specified key. {@code null} is not allowed.
37+
* @param ttl The value time-to-live, after which the value is automatically evicted.
38+
*/
39+
void put(K key, V value, Duration ttl)
40+
41+
}

0 commit comments

Comments
 (0)