New output buffer implementation

[dain]

No description provided.

[dain]

Note: Github, in its infinite wisdom, has shuffled the commit order.

[dain]

@haozhun can you review "Add new PartitionedBuffer" and "Add new BroadcastBuffer"

[haozhun]

ImmutableMap.of() will just work

[dain]

I renamed the classes to make it clear which ones are specific to the old SharedBuffer code and which ones are generic. I also named OutputBuffer implementations to end with "OutputBuffer" for clarity.

• Rename SharedBufferInfo to OutputBufferInfo
• Rename SharedBuffer to SharedOutputBuffer
• Rename SharedBufferMemoryManager to OutputBufferMemoryManager

I reworked the URI and BufferId management logic in the query scheduler. The key insight was that I could force the partition number and the taskId to be the same for partitioned outputs, with some minor changes. With that, I didn't need most of my changes, so the old code is restored:

• Guarantee task id and output partition are the same
• Remove unused partition from OutputBufferManager
• Remove unused partition from RemoteTask
• Declare out buffers immediately for non-broadcast stages

Finally, when working on the new BroadcastOutputBuffer, I discovered that I could reuse the internal PartitionedOutputBuffer.Partition object in BroadcastOutputBuffer by adding reference counting. I pulled out this class into a ClientBuffer and added tests for the client protocol. I'll squash the first two commits:

• Extract PartitionedOutputBuffer.Partition into ClientBuffer
• Add reference counting to ClientBuffer
• Add new BroadcastOutputBuffer

[haozhun]

• Looks good
  • Move SharedBuffer and related classes to new package
  • Add OutputBuffers.createInitialEmptyOutputBuffers
  • Rename SharedBufferInfo to OutputBufferInfo
  • Rename SharedBuffer to SharedOutputBuffer
  • Rename SharedBufferMemoryManager to OutputBufferMemoryManager
  • Fix thread safety issue in BroadcastOutputBufferManager
  • Change TaskId.id to be an int
  • Replace TaskId in output buffers with a generic OutputBufferId
  • Remove unused partition from OutputBufferManager
  • Remove unused partition from RemoteTask
  • Register output buffers before adding splits
  • Add LazyOutputBuffer to delay creation until plan is sent
  • Add buffer type to OutputBuffers to allow for different buffer implementations
  • Extract PartitionedOutputBuffer.Partition into ClientBuffer
• Extract interface OutputBuffer from SharedBuffer
  • Looks good, except comments below
  • OutputBuffer.setNoMorePages javadoc: ignored
• Guarantee task id and output partition are the same
  • Looks good
  • I would like to understand callers of SqlStageExecution.scheduleTask/Splits better
• Declare out buffers immediately for non-broadcast stages
  • Looks good, except minor comments below
  • PartitionedOutputBufferManager.addOutputBuffers
    • Validate that noMoreBuffer does not change from true to false, unless there's a valid use case to do that.
    • Calling withBuffer a bunch of times will allocate a bunch of objects and create the version a lot. I guess these don't really matter. But I would create a map and then pass the map in just to be on the safe side.
• Move shared buffer not-full notification to a new thread
  • SharedOutputBuffer: duplicate requireNonNull for maxBufferSize and systemMemoryUsageListener
  • Question about OutputBufferMemoryManager:
    • Is there any particular reason updateMemoryUsage submit the job outside the sync block and setNoBlockOnFull submit the job inside the sync block?
• Add new PartitionedOutputBuffer
  • Why is calling checkFlushComplete necessary/useful in setNoMorePages or the constructor?
• Add reference counting to ClientBuffer
  • In destroy, pendingRead should be completed with emptyResults(taskInstanceId, sequenceId, true)
  • In PageReference.addReference, checkState(referenceCount.getAndIncrement() == 0, "...")
  • I didn't know about FieldAccessNotGuarded, that's super nice. Let's add GuardedBy to currentSequenceId and destroyed
• Add new BroadcastOutputBuffer
  • enqueue not thread safe: gap before safeGetBuffersSnapshot can mean duplicate pages.
  • Discussion item: noMoreBuffers: put dereferencePage in finally. Other usages of dereferencePage may also benefit if they get put in a finally block. Should we do that? Or should we care at all?

Note: Please squash "Extract PartitionedOutputBuffer.Partition into ClientBuffer". For these multi-threaded code, I find it easier to reason about correctness when I have the full picture.

[dain]

* Declare out buffers immediately for non-broadcast stages
  * PartitionedOutputBufferManager.addOutputBuffers
    * Validate that noMoreBuffer does not change from `true` to `false`, unless there's a valid use case to do that.
    * Calling `withBuffer` a bunch of times will allocate a bunch of objects and create the `version` a lot. I guess these don't really matter. But I would create a map and then pass the map in just to be on the safe side.

The noMoreBuffer check is handled by the validate code, and the object allocation is not really a problem when you consider that Map will allocate an object for each entry anyway (also partitions is typically small).

The real issues is with the version number which the checkValidTransition code validates. I just rewrote the code to check the buffers directly and avoided the checkValidTransition code.

* Move shared buffer not-full notification to a new thread
  * Question about `OutputBufferMemoryManager`:
    * Is there any particular reason `updateMemoryUsage` submit the job outside the sync block and `setNoBlockOnFull` submit the job inside the sync block?

No, I simplified the code.

* Add new PartitionedOutputBuffer
  * Why is calling `checkFlushComplete` necessary/useful in `setNoMorePages` or the constructor?

setNoMorePages can free a reader which will see the finished flag. Then there is a race to destroy the buffers. It is simpler to just double check at the end of the method, then try to determine if this is a benign race. As for the constructor, it is likely not needed, but doesn't hurt anything (potentially, you could have zero partitions, but that seems weird).

* Add reference counting to ClientBuffer
  * In `destroy`, `pendingRead` should be completed with `emptyResults(taskInstanceId, sequenceId, true)`

This doesn't matter. There is only one client and the destroy comes from the client. This has the benefit of keeping the API surface area smaller.

[dain]

I will merge once the current release goes out