Doctrine: Related entities are not updated after changing id in EventSubscriber (onFlush) - orm

Question

I use composite primary keys for all Doctrine entities (multi-tenancy).
Example (yaml file):
id:
website:
associationKey: true
id:
type: integer
So I have to auto increment the id manually. All entities have id 1 as default.
Example (EventSubscriber):
class MultitenantEventSubscriber implements EventSubscriber
{
public function getSubscribedEvents()
{
return [
'onFlush'
];
}
public function onFlush(OnFlushEventArgs $event)
{
$em = $event->getEntityManager();
$uow = $em->getUnitOfWork();
foreach ($uow->getScheduledEntityInsertions() as $entity) {
if ($entity instanceof AutoIncrementInterface) {
$maxId = $em
->createQueryBuilder('e')
->select('MAX(e.id)')
->from(get_class($entity), 'e')
->getQuery()
->getSingleScalarResult();
if ($maxId > 1) {
$entity->setId($maxId + 1);
$md = $em->getClassMetadata(get_class($entity));
$uow->recomputeSingleEntityChangeSet($md, $entity);
}
}
}
}
}
Everything is working so far.
Problem:
Foreign key in related entities is not updated when I recompute the change set. Entity A with for example auto incremented id 23 is created correctly, but Entity B with reference to A still returns reference id 1 instead of id 23.
$entityA = new A();
$entityA->setValue('foo');
$entityB = new B();
$entityB->setValue('bar');
$entityB->setA($entityA);
$em->persist($entityA);
$em->persist($entityB);
$em->flush();
echo $entityA->getId(); // Returns 23
echo $entityB->getA()->getId() // Returns 1 instead of 23
Of course I could update the entityB in the EventSubscriber, but I need a generic solution because there are a lot of entities with different relations.
Maybe someone has an idea how to "recompute"/"refresh" related entities?