Wrong transaction as current transaction in datastore callbacks

[step76]

Hello,

I'm using the datastore callbacks to log the entity groups involved in a single transaction. The callbacks fire correctly but I've noticed in the logs some elements for which I explicitly passed null as a transaction.

So, digging the code I've found that the context passed to datastore callbacks is not populated with the transaction received in the get/put/delete/prepare methods, but is taken from an outer context.

  @Override
  public Future<Map<Key, Entity>> get(@Nullable Transaction txn, Iterable<Key> keys) {
    if (keys == null) {
      throw new NullPointerException("keys cannot be null");
    }

    // TODO: The handling of the Keys is pretty ugly.  We get an Iterable from the
    // user.  We need to copy this to a random access list (that we can mutate) for the PreGet
    // stuff. We will also convert it to a HashSet to support contains checks (for the RemoteApi
    // workaround).  There are also some O(N) calls to remove Keys from a List.
    List<Key> keyList = Lists.newArrayList(keys);

    // Allocate the Map that will receive the result of the RPC here so that PreGet callbacks can
    // add results.
    Map<Key, Entity> resultMap = new HashMap<Key, Entity>();
    PreGetContext preGetContext = new PreGetContext(this, keyList, resultMap);
    datastoreServiceConfig.getDatastoreCallbacks().executePreGetCallbacks(preGetContext);

    // Don't fetch anything from datastore that was provided by the preGet hooks.
    keyList.removeAll(resultMap.keySet());

    // Send the RPC(s).
    Future<Map<Key, Entity>> result = doBatchGet(txn, Sets.newLinkedHashSet(keyList), resultMap);

    // Invoke the user post-get callbacks.
    return new PostLoadFuture(result, datastoreServiceConfig.getDatastoreCallbacks(), this);
  }

For example, in the get method the context is created using this as transaction provider whilst the transaction is explicitly passed in the method, this lead to a wrong current transaction inside the callback.

Thank you
Stefano

[ludoch]

Asking @dgay42

[ludoch]

Filed internal issue b/281725939