Skip to content

Commit 83e9851

Browse files
ErmiasGErmiasG
authored
[HWORKS-527] fix jupyter jwt monitor for HA (#1444)
1 parent 42c330c commit 83e9851

File tree

12 files changed

+346
-49
lines changed

12 files changed

+346
-49
lines changed

‎hopsworks-common/src/main/java/io/hops/hopsworks/common/jupyter/CidAndPort.java

+3-1
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,11 @@
1616

1717
package io.hops.hopsworks.common.jupyter;
1818

19+
import java.io.Serializable;
1920
import java.util.Objects;
2021

21-
public class CidAndPort {
22+
public class CidAndPort implements Serializable {
23+
private static final long serialVersionUID = -7736027812979433344L;
2224
String cid;
2325
Integer port;
2426

‎hopsworks-common/src/main/java/io/hops/hopsworks/common/jupyter/JupyterJWT.java

+8
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,14 @@ public JupyterJWT(Project project, Users user, LocalDateTime expiration, CidAndP
3737
this.pidAndPort = pidAndPort;
3838
}
3939

40+
public JupyterJWT(Project project, Users user, LocalDateTime expiration, CidAndPort pidAndPort, String token,
41+
Path tokenFile) {
42+
super(project, user, expiration);
43+
this.pidAndPort = pidAndPort;
44+
this.tokenFile = tokenFile;
45+
this.token = token;
46+
}
47+
4048
@Override
4149
public boolean equals(Object o) {
4250
if (this == o) {

‎hopsworks-common/src/main/java/io/hops/hopsworks/common/jupyter/JupyterJWTCache.java

+128
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,128 @@
1+
/*
2+
* This file is part of Hopsworks
3+
* Copyright (C) 2023, Hopsworks AB. All rights reserved
4+
*
5+
* Hopsworks is free software: you can redistribute it and/or modify it under the terms of
6+
* the GNU Affero General Public License as published by the Free Software Foundation,
7+
* either version 3 of the License, or (at your option) any later version.
8+
*
9+
* Hopsworks is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY;
10+
* without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
11+
* PURPOSE. See the GNU Affero General Public License for more details.
12+
*
13+
* You should have received a copy of the GNU Affero General Public License along with this program.
14+
* If not, see <https://www.gnu.org/licenses/>.
15+
*/
16+
package io.hops.hopsworks.common.jupyter;
17+
18+
import com.hazelcast.core.HazelcastInstance;
19+
import com.hazelcast.map.IMap;
20+
import com.hazelcast.query.Predicate;
21+
import com.hazelcast.query.Predicates;
22+
import io.hops.hopsworks.common.util.DateUtils;
23+
24+
import javax.ejb.Singleton;
25+
import javax.ejb.TransactionAttribute;
26+
import javax.ejb.TransactionAttributeType;
27+
import javax.inject.Inject;
28+
import java.nio.file.Paths;
29+
import java.util.Collection;
30+
import java.util.HashMap;
31+
import java.util.Iterator;
32+
import java.util.Optional;
33+
import java.util.Set;
34+
import java.util.TreeSet;
35+
36+
@Singleton
37+
@TransactionAttribute(TransactionAttributeType.NEVER)
38+
public class JupyterJWTCache {
39+
private static final String MAP_NAME = "jupyterJWTMap";
40+
@Inject
41+
private HazelcastInstance hazelcastInstance;
42+
43+
private final TreeSet<JupyterJWTDTO> jupyterJWTs = new TreeSet<>((t0, t1) -> {
44+
if (t0.equals(t1)) {
45+
return 0;
46+
} else {
47+
if (t0.getExpiration().isBefore(t1.getExpiration())) {
48+
return -1;
49+
} else if (t0.getExpiration().isAfter(t1.getExpiration())) {
50+
return 1;
51+
}
52+
return 0;
53+
}
54+
});
55+
56+
private final HashMap<CidAndPort, JupyterJWT> pidAndPortToJWT = new HashMap<>();
57+
58+
public void add(JupyterJWT jupyterJWT) {
59+
if (hazelcastInstance != null) {
60+
IMap<CidAndPort, JupyterJWTDTO> pidAndPortToJWTMap = hazelcastInstance.getMap(MAP_NAME);
61+
pidAndPortToJWTMap.put(jupyterJWT.pidAndPort, new JupyterJWTDTO(jupyterJWT));
62+
} else {
63+
jupyterJWTs.add(new JupyterJWTDTO(jupyterJWT));
64+
pidAndPortToJWT.put(jupyterJWT.pidAndPort, jupyterJWT);
65+
}
66+
}
67+
68+
public Optional<JupyterJWT> get(CidAndPort pidAndPort) {
69+
if (hazelcastInstance != null) {
70+
IMap<CidAndPort, JupyterJWTDTO> pidAndPortToJWTMap = hazelcastInstance.getMap(MAP_NAME);
71+
JupyterJWTDTO jupyterJWTDTO = pidAndPortToJWTMap.get(pidAndPort);
72+
if (jupyterJWTDTO != null) {
73+
return Optional.of(
74+
new JupyterJWT(jupyterJWTDTO.getProject(), jupyterJWTDTO.getUser(), jupyterJWTDTO.getExpiration(),
75+
pidAndPort, jupyterJWTDTO.getToken(), Paths.get(jupyterJWTDTO.getTokenFile())));
76+
}
77+
return Optional.empty();
78+
} else {
79+
return Optional.ofNullable(pidAndPortToJWT.get(pidAndPort));
80+
}
81+
}
82+
83+
public void remove(CidAndPort pidAndPort) {
84+
if (hazelcastInstance != null) {
85+
IMap<CidAndPort, JupyterJWTDTO> pidAndPortToJWTMap = hazelcastInstance.getMap(MAP_NAME);
86+
pidAndPortToJWTMap.remove(pidAndPort);
87+
} else {
88+
JupyterJWT jupyterJWT = pidAndPortToJWT.remove(pidAndPort);
89+
jupyterJWTs.remove(new JupyterJWTDTO(jupyterJWT));
90+
}
91+
}
92+
93+
public void replaceAll(Set<JupyterJWT> renewedJWTs) {
94+
if (hazelcastInstance != null) {
95+
IMap<CidAndPort, JupyterJWTDTO> pidAndPortToJWTMap = hazelcastInstance.getMap(MAP_NAME);
96+
renewedJWTs.forEach(t -> pidAndPortToJWTMap.replace(t.pidAndPort, new JupyterJWTDTO(t)));
97+
} else {
98+
renewedJWTs.forEach(t -> {
99+
//remove old token
100+
JupyterJWT jupyterJWT = pidAndPortToJWT.remove(t.pidAndPort);
101+
jupyterJWTs.remove(new JupyterJWTDTO(jupyterJWT));
102+
//Add the new token
103+
jupyterJWTs.add(new JupyterJWTDTO(t));
104+
pidAndPortToJWT.put(t.pidAndPort, t);
105+
});
106+
}
107+
}
108+
109+
public int getSize() {
110+
if (hazelcastInstance != null) {
111+
IMap<CidAndPort, JupyterJWTDTO> pidAndPortToJWTMap = hazelcastInstance.getMap(MAP_NAME);
112+
return pidAndPortToJWTMap.size();
113+
} else {
114+
return jupyterJWTs.size();
115+
}
116+
}
117+
118+
public Iterator<JupyterJWTDTO> getMaybeExpired() {
119+
if (hazelcastInstance != null) {
120+
IMap<CidAndPort, JupyterJWTDTO> pidAndPortToJWTMap = hazelcastInstance.getMap(MAP_NAME);
121+
Predicate<CidAndPort, JupyterJWTDTO> expirationPredicate = Predicates.lessEqual("expiration", DateUtils.getNow());
122+
Collection<JupyterJWTDTO> jupyterJWTDTOS = pidAndPortToJWTMap.values(expirationPredicate);
123+
return jupyterJWTDTOS.iterator();
124+
} else {
125+
return jupyterJWTs.iterator();
126+
}
127+
}
128+
}

‎hopsworks-common/src/main/java/io/hops/hopsworks/common/jupyter/JupyterJWTDTO.java

+113
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,113 @@
1+
/*
2+
* This file is part of Hopsworks
3+
* Copyright (C) 2023, Hopsworks AB. All rights reserved
4+
*
5+
* Hopsworks is free software: you can redistribute it and/or modify it under the terms of
6+
* the GNU Affero General Public License as published by the Free Software Foundation,
7+
* either version 3 of the License, or (at your option) any later version.
8+
*
9+
* Hopsworks is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY;
10+
* without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
11+
* PURPOSE. See the GNU Affero General Public License for more details.
12+
*
13+
* You should have received a copy of the GNU Affero General Public License along with this program.
14+
* If not, see <https://www.gnu.org/licenses/>.
15+
*/
16+
package io.hops.hopsworks.common.jupyter;
17+
18+
import io.hops.hopsworks.common.util.DateUtils;
19+
import io.hops.hopsworks.persistence.entity.project.Project;
20+
import io.hops.hopsworks.persistence.entity.user.Users;
21+
22+
import java.io.Serializable;
23+
import java.time.LocalDateTime;
24+
import java.util.Objects;
25+
26+
public class JupyterJWTDTO implements Serializable {
27+
private static final long serialVersionUID = -5687462769985361531L;
28+
private Project project;
29+
private Users user;
30+
private LocalDateTime expiration;
31+
private String token;
32+
private String tokenFile;
33+
private final CidAndPort pidAndPort;
34+
35+
public JupyterJWTDTO(JupyterJWT jupyterJWT) {
36+
this.project = jupyterJWT.project;
37+
this.user = jupyterJWT.user;
38+
this.expiration = jupyterJWT.expiration;
39+
this.token = jupyterJWT.token;
40+
this.tokenFile = jupyterJWT.tokenFile.toString();
41+
this.pidAndPort = jupyterJWT.pidAndPort;
42+
}
43+
44+
public Project getProject() {
45+
return project;
46+
}
47+
48+
public void setProject(Project project) {
49+
this.project = project;
50+
}
51+
52+
public Users getUser() {
53+
return user;
54+
}
55+
56+
public void setUser(Users user) {
57+
this.user = user;
58+
}
59+
60+
public LocalDateTime getExpiration() {
61+
return expiration;
62+
}
63+
64+
public void setExpiration(LocalDateTime expiration) {
65+
this.expiration = expiration;
66+
}
67+
68+
public String getToken() {
69+
return token;
70+
}
71+
72+
public void setToken(String token) {
73+
this.token = token;
74+
}
75+
76+
public String getTokenFile() {
77+
return tokenFile;
78+
}
79+
80+
public void setTokenFile(String tokenFile) {
81+
this.tokenFile = tokenFile;
82+
}
83+
84+
public CidAndPort getPidAndPort() {
85+
return pidAndPort;
86+
}
87+
88+
public boolean maybeRenew(LocalDateTime now) {
89+
return now.isAfter(expiration) || now.isEqual(expiration);
90+
}
91+
92+
public boolean isExpired() {
93+
LocalDateTime now = DateUtils.getNow();
94+
return now.isAfter(expiration) || now.isEqual(expiration);
95+
}
96+
97+
@Override
98+
public boolean equals(Object o) {
99+
if (this == o) {
100+
return true;
101+
}
102+
if (o == null || getClass() != o.getClass()) {
103+
return false;
104+
}
105+
JupyterJWTDTO that = (JupyterJWTDTO) o;
106+
return Objects.equals(project.getId(), that.project.getId()) && Objects.equals(user.getUid(), that.user.getUid());
107+
}
108+
109+
@Override
110+
public int hashCode() {
111+
return Objects.hash(project.getId(), user.getUid());
112+
}
113+
}

0 commit comments

Comments
 (0)