addresses slow import/export performance by limiting persistence context size - #37926

Also, @pedroigor I don't quite follow the existing logic that seems to be attempting something similar. Based upon this we're skipping all merge/persist calls

I always thought that the import operation was mainly targeted for provisioning realms with the very basic data and anything beyond that would be better addressed by dumping/restoring the database. But reading the original it seems we introduced a fix that is degrading the performance of import, is that correct?

I see the solution you are proposing but for me, it is adding additional complexity for use cases that I'm not sure we are supposed to support. It looks like we are trying to work around the problem that we now have the import running in the scope of a single (KC/HHM) session.

Wouldn't be better to import user files in batch (e.g.: org.keycloak.exportimport.util.ExportImportSessionTask#useExistingSession returning false) instead? Or even if using a single session, a mechanism that allows to configure how to process entries in batches?

But reading the original it seems we introduced a fix that is degrading the performance of import, is that correct?

Import from multiple user files, yes.

I see the solution you are proposing but for me, it is adding additional complexity for use cases that I'm not sure we are supposed to support. It looks like we are trying to work around the problem that we now have the import running in the scope of a single (KC/HHM) session.

I don't agree really agree with this viewpoint - the existing logic is inefficient for grouping together any number of user creations. I'd rather start with addressing that, then decide if splitting over multiple transactions makes sense.

Wouldn't be better to import user files in batch (e.g.: org.keycloak.exportimport.util.ExportImportSessionTask#useExistingSession returning false) instead?

Did you see the rationale for why this was set to true?

Or even if using a single session, a mechanism that allows to configure how to process entries in batches?

That's what this code change will effectively allow for - by removing the explicit and implicit flushes we'll allow hibernate to actually batch under the covers.

Dumping/restoring, but also migrating from a legacy system? Import is also for this use case no?

I see the solution you are proposing but for me, it is adding additional complexity for use cases that I'm not sure we are supposed to support. It looks like we are trying to work around the problem that we now have the import running in the scope of a single (KC/HHM) session.

I don't agree really agree with this viewpoint - the existing logic is inefficient for grouping together any number of user creations. I'd rather start with addressing that, then decide if splitting over multiple transactions makes sense.

I think it does because a single transaction might eventually end up with lot of objects in the persistence context and that, I think, is the limiting factor here.

At the same time, it would be nice to improve how batch processing from PersistenceExceptionConverter works when flush is explicitly invoked. Like when creating users.

Wouldn't be better to import user files in batch (e.g.: org.keycloak.exportimport.util.ExportImportSessionTask#useExistingSession returning false) instead?

Did you see the rationale for why this was set to true?

Yes, I see. Do you mean what you commented here #34364 (comment)?

Or even if using a single session, a mechanism that allows to configure how to process entries in batches?

That's what this code change will effectively allow for - by removing the explicit and implicit flushes we'll allow hibernate to actually batch under the covers.

IIRC, most of the time (mainly for users) we are calling flush to be able to handle any constraint from the database (not sure if the best practice) before finishing the response. Removing the explicit flush would be nice and I remember discussions about this in the past. Others might remember the actual implications and why we did not change that until now. In addition to user creation, I think we also have flush being called in other places like when joining and leaving groups.

In summary, I think we agree/need:

• Prevent the persistence context from growing and growing when processing entries in a single unit of work.
• Improve the batch mechanism we already have in the PersistenceExceptionConverter

The only thing I'm not sure is reviewing the explicit calls to flush. That will be tuff to review/remove ...

The only thing I'm not sure is reviewing the explicit calls to flush. That will be tuff to review/remove ...

Thank you for the review @pedroigor and agreed on this point. To make things clearer / safer, I've refined the changes to be more explicit rather than to hide all the handling under the PersistenceExceptionConverter. I also stuck with using a thread local so that the batching can apply to all entitymanagers proxied by PersistenceExceptionConverter - but that does make the JpaConnectionProvider.batchMode method a bit odd because it's effectively static.

Prevent the persistence context from growing and growing when processing entries in a single unit of work.

To keep things explicit it's now up to a caller to initiate a batching mode, then turn it off. Then to align with the previous import logic this is done around each directory file. Alternatively the notion of how many operations to batch could be handled by the PersistenceExceptionConverter like the migration logic.

This change by itself gives us most of the performance benefit of using separate transactions. There are 2 further potential optimizations:

1. Changing the query flush modes to COMMIT via PersistenceExceptionConverter when in batching mode. I know this is not ideal and potentially not generally correct - but you can definitely see, especially with larger persistence contexts, that determining when to flush is taking up a lot of time. With limiting the persistence context size this may no longer be needed and/or this could be done in a targeted way - setting the flush mode to COMMIT where appropriate on where the queries are used.
2. Check the PersistenceExceptionConverter.isBatchMode before doing explicit flushes or detaches which futher allows for true batching of the underlying operations by hibernate.

Improve the batch mechanism we already have in the PersistenceExceptionConverter

I left this out of the changes for now, but yes that would definitely be nice.

Some small benchmarking efforts with running pg in podman and 2000 minimal users in a single file:

• no optimizations: 28 s
• without batching (not inhibiting flush / detach), but with setting queries to COMMIT: 22 s
• with batching, but without setting queries to COMMIT: 15.5 s
• with batching and with setting queries to COMMIT: 1.2 s

10000 users in a single file with batching and with setting queries to COMMIT took 4 s, so pretty linear time scaling. If we can agree on an implicit flush behavior (like the current logic for migration batching), then we wouldn't have to scope the batching behavior per file nor worry about the persistence context getting too large.

Looked at #34364 (comment) - I don't see any large performance issue even with the current state of main as long as the users per file is under 1000, so there doesn't seem to be a regression with in general. However doing several thousand users per file starts to exhibit the same type of performance issues as importing a large number of users.

I looked at the explicit flushes in the past, and IMHO the primary reason they are there is the fear what might happen if we remove them. Even if it has been added in the past to show DB constrains earlier and not only when the transaction is over, this is IMHO an antipattern as with concurrent requests there could still be conflicts discovered only at commit time.

So +1 for for removing them step-by-step. This might surface some wrongly applied JPA logic especially when child entities are updated - we've been struck by #11666 in the past, which is AFAIK still a problem in Hibernate.

Also +1 for the COMMIT setting when importing.

When importing a list of users, IMHO it still makes sense to put each user file in its own transaction to avoid allocating too much memory and having a large persistence context.

@shawkins Nice work!

JFYI - not batched flushes are also an issue for default client scopes when the realm is created. So there might also be some possible improvement.

When importing a list of users, IMHO it still makes sense to put each user file in its own transaction to avoid allocating too much memory and having a large persistence context.

I'll show what a possible higher-level control for this could look like in my next commit. I believe that something like this could eventually be used to allow callers to do error handling prior to the commit.

JFYI - not batched flushes are also an issue for default client scopes when the realm is created. So there might also be some possible improvement.

Thanks @mabartos that's good to know, I can certainly expand the scope of the batching as needed.

I'll show what a possible higher-level control for this could look like in my next commit.

This commit adds batching to export, so that the export of 10000 users to a single file or the realm file is taking around 14-18 s - so there is some asymetric inefficiency as import from a single file only takes about 4 seconds. Since #37990 has not been applied, these results are with running in a single transaction.

The higher level batch controls provide flush / clear methods when running in batch mode. Usage of it is shown in two ways - for export it's just clearing the context around what could be transaction boundaries. For import it's flushing and clearing the context after a given number of users, rather than just letting the number of users in the file drive the decision. I'm open to tweaking this in anyway that others see fit. In the end we should be able to batch operations and limit the memory usage of the persistence context in a flexible way.

Also by leaving the ExportImportSessionTask my intention is to allow a fallback to multi-transaction behavior, but to make that require an additional switch for import / export.

Adding hold label not to merge before 26.2 release.

Refactored to pull out unrelated changes, command level switches to use multiple transactions, and the KeycloakModelUtils handling of reusing the current session.

Also switched the logic to create nested entity managers when batching.

I don't believe there is any benefit to exposing to the user multi-txn controls - however I still left in the current structuring for multi-txn in case we need to quickly support that.

@mabartos also added batching for the default client scopes - as you pointed out they are the next biggest issue with import performance. And it does are to aslo be hit when migration logic runs - along with poor query performance in the migrateRealm method. Those are the likely culprits behind #34662

@vmuzikar @ahus1 @pedroigor what needs to be done here to keep this moving forward? I believe the changes have been minimized to keep this lower risk, but I'm open to refining as needed.

I plan to review on Monday. Additionally, I'd like to ask if we could get a review either from @pedroigor or @ahus1.

After that, IMHO we'll be ready to merge.

@shawkins - thank you for the PR. While I suppose this is good to be merged, see below for some thoughts about corner cases.

Note that this is built on the assumption that importing a user doesn't modify any other existing entity that has been touched in a previous batch. Those entities are basically detached by the em.clear() method. Those now detached entities are still bound to the current session in UserCacheSession#managedUsers. Any subsequent changes to them will get lost.

Operations like https://github.com/keycloak/keycloak/blob/9721d920feeed1718bda49ac59cb4fc80ae305a2/model/storage-services/src/main/java/org/keycloak/exportimport/util/ImportUtils.java#L73-L78 still seems to only add new entities, but don't modify existing ones.

I'm good to proceed with this as long as we don't extend this to other areas of Keycloak where this assumption might not hold. While I am quite sure this will not break Keycloak code, I hope people didn't write extensions that would break this assumption.

While the speed of the import is probably good, with a lot of users or other entities it will acquire a lot of entries in memory due to UserCacheSession#managedUsers and other similar patterns I suspect for other entities. I see that the recommendation to put users in files is now removed completely. Assuming that the files are still imported in separate transactions, and that this would keep the heap usage at bay, we might still want to keep that recommendation, with a higher number though.

Note that this is built on the assumption that importing a user doesn't modify any other existing entity that has been touched in a previous batch.

Which is an assumption that was already built into the multi-file/transaction case, so this is not a new concept. Here we're mostly just replacing the transaction with a manipulation of the EntityManager(s).

While I am quite sure this will not break Keycloak code, I hope people didn't write extensions that would break this assumption.

Since the EntityManagers class is in spi-private, they would have to knowingly abuse a private api.

Assuming that the files are still imported in separate transactions

With this change they are not.

and that this would keep the heap usage at bay, we might still want to keep that recommendation, with a higher number though.

Right with the old text gone it doesn't cover the possiblity of the import running out of memory. I add some language about using multiple files to protect the memory of the import operation. I did see that without any tweaks 200,000 bare-bones users were enough to run the import process out of memory.

As for what the number should be:

• just leave it at 50, since it doesn't make much performance difference - addresses slow import/export performance by limiting persistence context size #37926 (comment) - and should we need to have a mode of operation that still utilizes multiple transactions this would safe
• A new default, such as 1000 / file.

While I am quite sure this will not break Keycloak code, I hope people didn't write extensions that would break this assumption.

Since the EntityManagers class is in spi-private, they would have to knowingly abuse a private api.

I had more something like a very crazy user storage provider in mind or similar. Anyway, it should be fine.

Assuming that the files are still imported in separate transactions

With this change they are not.

So we are now limiting the import to the heap available to KC during import? If you think this is acceptable, it might be worth mentioning in the docs ("if you see an OOM during import due to a large number of imported entities, increase the heap available to Keycloak"). By running each file in a separate transaction, the limitation by the available heap can possibly be avoided.

Given your tests, the new default number of users per file can be increased. Given you latest comment, I suppose it would be ok to increase it to 100_000 users then?

So we are now limiting the import to the heap available to KC during import?

This is not a new limitation. This has been a latent issue that was probably not addressed due to recommending multiple files.

The old logic put the entire user contents from a single file (whether that was realm+users, or just a users file) into memory.

That is still the case, even with these changes, for a file containg both the realm and the users.

You probably noticed that I tried to improve a little upon multi-file import logic with some tweaks so that it could be more streaming, but that problem wasn't expressly my goal to address here - just the multi-transaction issue.

So even with this pr we still do eventually hit an OOM at some number of users per file.

By running each file in a separate transaction, the limitation by the available heap can possibly be avoided.

No, that is central point of these changes - multiple transactions are not needed.

However the usage of multiple files still helps address some other limitation with our logic and how it is parsing / performing the import.

With a little more effort this limitation could be removed.

The changes here will allow me to address #38251 by making the logic more uniform between single and multi-file imports, then I believe I could tackle better streaming of user imports such that multiple files are no longer even needed.

@shawkins - ok, let's then process with this one as it is a big improvement to many, and tackle limitations in future PRs.

@ahus1 sounds good. Refined the language in the docs to reference memory limitations.

@shawkins @vmuzikar Wouldn't be good to brag a little bit about these improvements in the 26.3.0 release notes? :P I'd say yes.

@shawkins @vmuzikar Wouldn't be good to brag a little bit about these improvements in the 26.3.0 release notes? :P I'd say yes.

I'll add something on to #39377