Group persistent session work activities in parent span or link them

Closes #35430

Signed-off-by: Alexander Schwartz <aschwart@redhat.com>
This commit is contained in:
Alexander Schwartz
2024-12-10 13:55:35 +01:00
committed by GitHub
parent f21bbb26d5
commit cde8f25cc2
3 changed files with 68 additions and 33 deletions

View File

@@ -191,6 +191,9 @@ abstract public class PersistentSessionsChangelogBasedTransaction<K, V extends S
List<SessionChangesPerformer<K, V>> changesPerformers = null;
for (Map.Entry<K, SessionUpdatesList<V>> entry : Stream.concat(updates.entrySet().stream(), offlineUpdates.entrySet().stream()).toList()) {
SessionUpdatesList<V> sessionUpdates = entry.getValue();
if (sessionUpdates.getUpdateTasks().isEmpty()) {
continue;
}
SessionEntityWrapper<V> sessionWrapper = sessionUpdates.getEntityWrapper();
V entity = sessionWrapper.getEntity();
boolean isOffline = entity.isOffline();

View File

@@ -17,14 +17,20 @@
package org.keycloak.models.sessions.infinispan.changes;
import io.opentelemetry.api.trace.Span;
import io.opentelemetry.api.trace.SpanBuilder;
import io.opentelemetry.api.trace.Tracer;
import io.opentelemetry.context.Context;
import org.jboss.logging.Logger;
import org.keycloak.common.util.Retry;
import org.keycloak.models.KeycloakSessionFactory;
import org.keycloak.models.utils.KeycloakModelUtils;
import org.keycloak.tracing.TracingProvider;
import java.time.Duration;
import java.time.temporal.ChronoUnit;
import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.TimeUnit;
@@ -81,41 +87,60 @@ public class PersistentSessionsWorker {
if (polled != null) {
batch.add(polled);
queue.drainTo(batch, maxBatchSize - 1);
try {
LOG.debugf("Processing %d deferred session updates.", batch.size());
Retry.executeWithBackoff(iteration -> {
if (iteration < 2) {
// attempt to write whole batch in the first two attempts
KeycloakModelUtils.runJobInTransaction(factory,
innerSession -> batch.forEach(c -> c.perform(innerSession)));
batch.forEach(PersistentUpdate::complete);
} else {
LOG.warnf("Running single changes in iteration %d for %d entries", iteration, batch.size());
ArrayList<PersistentUpdate> performedChanges = new ArrayList<>();
List<Throwable> throwables = new ArrayList<>();
batch.forEach(change -> {
try {
KeycloakModelUtils.runJobInTransaction(factory,
change::perform);
change.complete();
performedChanges.add(change);
} catch (Throwable ex) {
throwables.add(ex);
KeycloakModelUtils.runJobInTransaction(factory, outerSession -> {
TracingProvider tracing = outerSession.getProvider(TracingProvider.class);
Tracer process = tracing.getTracer("PersistentSessionsWorker");
SpanBuilder spanBuilder = process.spanBuilder("PersistentSessionsWorker.process");
batch.stream().map(update -> update.getSpan().getSpanContext()).forEach(spanBuilder::addLink);
Span span = tracing.startSpan(spanBuilder);
List<Span> batchSpans = new LinkedList<>();
try {
batch.forEach(persistentUpdate -> {
// This adds another span to the parent span to avoid updating span links after span creation as suggested by the API
SpanBuilder sb = process.spanBuilder("PersistentSessionsWorker.batch");
sb.setParent(Context.current().with(persistentUpdate.getSpan()));
sb.addLink(span.getSpanContext());
batchSpans.add(sb.startSpan());
});
LOG.debugf("Processing %d deferred session updates.", batch.size());
Retry.executeWithBackoff(iteration -> {
if (iteration < 2) {
// attempt to write whole batch in the first two attempts
KeycloakModelUtils.runJobInTransaction(factory,
innerSession -> batch.forEach(c -> c.perform(innerSession)));
batch.forEach(PersistentUpdate::complete);
} else {
LOG.warnf("Running single changes in iteration %d for %d entries", iteration, batch.size());
ArrayList<PersistentUpdate> performedChanges = new ArrayList<>();
List<Throwable> throwables = new ArrayList<>();
batch.forEach(change -> {
try {
KeycloakModelUtils.runJobInTransaction(factory,
change::perform);
change.complete();
performedChanges.add(change);
} catch (Throwable ex) {
throwables.add(ex);
}
});
batch.removeAll(performedChanges);
if (!throwables.isEmpty()) {
RuntimeException ex = new RuntimeException("unable to complete some changes");
throwables.forEach(ex::addSuppressed);
throw ex;
}
});
batch.removeAll(performedChanges);
if (!throwables.isEmpty()) {
RuntimeException ex = new RuntimeException("unable to complete some changes");
throwables.forEach(ex::addSuppressed);
throw ex;
}
}
},
Duration.of(10, ChronoUnit.SECONDS), 0);
} catch (RuntimeException ex) {
batch.forEach(o -> o.fail(ex));
LOG.warnf(ex, "Unable to write %d deferred session updates", batch.size());
}
},
Duration.of(10, ChronoUnit.SECONDS), 0);
} catch (RuntimeException ex) {
tracing.error(ex);
batch.forEach(o -> o.fail(ex));
LOG.warnf(ex, "Unable to write %d deferred session updates", batch.size());
} finally {
batchSpans.forEach(Span::end);
tracing.endSpan();
}
});
}
}
}

View File

@@ -17,6 +17,7 @@
package org.keycloak.models.sessions.infinispan.changes;
import io.opentelemetry.api.trace.Span;
import org.keycloak.models.KeycloakSession;
import java.util.concurrent.CompletableFuture;
@@ -31,9 +32,11 @@ public class PersistentUpdate {
private final Consumer<KeycloakSession> task;
private final CompletableFuture<Void> future = new CompletableFuture<>();
private final Span span;
public PersistentUpdate(Consumer<KeycloakSession> task) {
this.task = task;
this.span = Span.current();
}
public void perform(KeycloakSession session) {
@@ -51,4 +54,8 @@ public class PersistentUpdate {
public CompletableFuture<Void> future() {
return future;
}
public Span getSpan() {
return span;
}
}