Lost FOR KEY SHARE lock

Hi,

While looking at another issue [1], I noticed that there's a race
condition in heap_update() where a FOR KEY SHARE lock can be lost.
Leaving out a lot of details, heap_update() looks like this:

/*
* If the tuple we're updating is locked, we need to preserve the locking
* info in the old tuple's Xmax. Prepare a new Xmax value for this.
*/
compute_new_xmax_infomask(HeapTupleHeaderGetRawXmax(oldtup.t_data),
oldtup.t_data->t_infomask,
oldtup.t_data->t_infomask2,
xid, *lockmode, true,
&xmax_old_tuple, &infomask_old_tuple,
&infomask2_old_tuple);
...
if (need_toast || newtupsize > pagefree)
{
...
/*
* To prevent concurrent sessions from updating the tuple, we have to
* temporarily mark it locked, while we release the page-level lock.
*
* To satisfy the rule that any xid potentially appearing in a buffer
* written out to disk, we unfortunately have to WAL log this
* temporary modification. We can reuse xl_heap_lock for this
* purpose. If we crash/error before following through with the
* actual update, xmax will be of an aborted transaction, allowing
* other sessions to proceed.
*/

/*
* Compute xmax / infomask appropriate for locking the tuple. This has
* to be done separately from the combo that's going to be used for
* updating, because the potentially created multixact would otherwise
* be wrong.
*/
compute_new_xmax_infomask(HeapTupleHeaderGetRawXmax(oldtup.t_data),
oldtup.t_data->t_infomask,
oldtup.t_data->t_infomask2,
xid, *lockmode, false,
&xmax_lock_old_tuple, &infomask_lock_old_tuple,
&infomask2_lock_old_tuple);

/* ... update the old tuple with the computed values ... */

LockBuffer(buffer, BUFFER_LOCK_UNLOCK);

/* ... toast or find new page ... */

/* Re-acquire the lock on the old tuple's page. */
LockBuffer(buffer, BUFFER_LOCK_EXCLUSIVE);
}

/* Clear obsolete visibility flags, possibly set by ourselves above... */
oldtup.t_data->t_infomask &= ~(HEAP_XMAX_BITS | HEAP_MOVED);
oldtup.t_data->t_infomask2 &= ~HEAP_KEYS_UPDATED;
/* ... and store info about transaction updating this tuple */
Assert(TransactionIdIsValid(xmax_old_tuple));
HeapTupleHeaderSetXmax(oldtup.t_data, xmax_old_tuple);
oldtup.t_data->t_infomask |= infomask_old_tuple;
oldtup.t_data->t_infomask2 |= infomask2_old_tuple;
HeapTupleHeaderSetCmax(oldtup.t_data, cid, iscombo);

While we're not holding the lock, another backend can acquire a KEY
SHARE lock on the tuple. But the 'xmax_old_tuple', 'infomask_old_tuple',
and 'infomask2_old_tuple' values were calculated before releasing the
lock, and won't include the new locker. By updating the old tuple with
those values, any KEY SHARE lock acquired in that window is lost.

Attached is an isolation test case to demonstrate that. I ran out of
energy trying to fix it; it doesn't seem that hard at first glance but
it's a little subtle. If you just recompute the values after locking the
tuple, then the values will also include the transaction's own lock on
the tuple. That's not wrong, but I think it will lead to a lot of extra
multixid. I think this needs to the members of the multixid, remove
itself as a locker, and add back as an updater.

Another observation is that if the transaction already has a suitable
lock on the tuple, there's no need to lock it again here. That would
seem like an easy optimization. heap_lock_tuple() already contains logic
for that, it could perhaps be reused here.

- Heikki