What type of PR is this?

Uncomment only one /kind <> line, hit enter to put that in a new line, and remove leading whitespaces from that line:

/kind feature

What does this PR do / why do we need it:

This PR allows our memory tracking context to keep track of the memory used by snapshot when the underlying instance is a SingleInputSnapshotState. This will help prevent any hanging caused by memory overuse as well as prevent Java OOM errors.

Which issue(s) this PR fixes:

Fixes #I4GM3X:Checkpoint - Improve snapshot memory management

Special notes for your reviewers: