You are right about race condition, but onReturn does not modify order at all

This is what happens in the background:

 - the `commerce_payment.checkout.return` route on commerce payment
   module is called which at first it let order payment gateway plugin
   to create an payment for order
 - then no matter what happens in the onReturn method of plugin order checkout flow will be redirected on another stage (if onReturn throws exception it will be redirected to previous stage else it will be redirect to next stage)
 - in this situation the method  `redirectToStep` in the
   `Drupal\commerce_checkout\Plugin\Commerce\CheckoutFlow\CheckoutFlowBase`
   will modify and save order which is where the problem happens.

this is the implementation of `redirectToStep`:

    /**
     * {@inheritdoc}
     */
    public function redirectToStep($step_id) {
      if (!$this->isStepVisible($step_id)) {
        throw new \InvalidArgumentException(sprintf('Invalid step ID "%s" passed to redirectToStep().', $step_id));
      }

      $this->order->set('checkout_step', $step_id);
      $this->onStepChange($step_id);
      $this->order->save();

      throw new NeedsRedirectException(Url::fromRoute('commerce_checkout.form', [
        'commerce_order' => $this->order->id(),
        'step' => $step_id,
      ])->toString());
    }

so to solve the problem in your onReturn method you should stand for onNotify changes to be done (a while loop which waits until changes on order are set)

hopes it helps.