Fix potential use-after-move in `RecvOp`.

Note that has not been a problem in practice because the callback
would only be used-after-move in an error condition
(`Rendezvous::ParseKey()` failing) that would only occur if a
malformed graph was fed to a TensorFlow server.

PiperOrigin-RevId: 176434460

1 files changed, 20 insertions, 13 deletions

diff --git a/tensorflow/core/kernels/sendrecv_ops.cc b/tensorflow/core/kernels/sendrecv_ops.cc
index 542382872c..206fd40fa6 100644
--- a/tensorflow/core/kernels/sendrecv_ops.cc
+++ b/tensorflow/core/kernels/sendrecv_ops.cc
@@ -142,17 +142,12 @@ RecvOp::RecvOp(OpKernelConstruction* ctx) : AsyncOpKernel(ctx) {
   }
 }
 
-void RecvOp::ComputeAsync(OpKernelContext* ctx, DoneCallback done) {
-  OP_REQUIRES(
-      ctx, ctx->rendezvous() != nullptr,
-      errors::Internal("Op kernel context needs to provide a rendezvous."));
-
-  Rendezvous::Args args;
-  args.device_context = ctx->op_device_context();
-  args.alloc_attrs = ctx->output_alloc_attr(0);
+namespace {
+Rendezvous::DoneCallback make_recv_callback(OpKernelContext* ctx,
+                                            AsyncOpKernel::DoneCallback done) {
   using namespace std::placeholders;
-  Rendezvous::DoneCallback done_cb = std::bind(
-      [ctx](DoneCallback done,
+  return std::bind(
+      [ctx](AsyncOpKernel::DoneCallback done,
             // Begin unbound arguments.
             const Status& s, const Rendezvous::Args& send_args,
             const Rendezvous::Args& recv_args, const Tensor& val,
@@ -170,19 +165,31 @@ void RecvOp::ComputeAsync(OpKernelContext* ctx, DoneCallback done) {
         done();
       },
       std::move(done), _1, _2, _3, _4, _5);
+}
+}  // namespace
+
+void RecvOp::ComputeAsync(OpKernelContext* ctx, DoneCallback done) {
+  OP_REQUIRES(
+      ctx, ctx->rendezvous() != nullptr,
+      errors::Internal("Op kernel context needs to provide a rendezvous."));
+
+  Rendezvous::Args args;
+  args.device_context = ctx->op_device_context();
+  args.alloc_attrs = ctx->output_alloc_attr(0);
 
   FrameAndIter frame_iter = GetFrameAndIter(ctx, hostmem_sendrecv_);
   if (frame_iter == FrameAndIter(0, 0)) {
     VLOG(2) << "Recv " << parsed_key_.buf_;
-    ctx->rendezvous()->RecvAsync(parsed_key_, args, std::move(done_cb));
+    ctx->rendezvous()->RecvAsync(parsed_key_, args,
+                                 make_recv_callback(ctx, std::move(done)));
   } else {
     Rendezvous::ParsedKey in_loop_parsed;
     GetRendezvousKey(key_prefix_, frame_iter, &in_loop_parsed.buf_);
     VLOG(2) << "Recv " << in_loop_parsed.buf_;
     OP_REQUIRES_OK_ASYNC(
         ctx, Rendezvous::ParseKey(in_loop_parsed.buf_, &in_loop_parsed), done);
-
-    ctx->rendezvous()->RecvAsync(in_loop_parsed, args, std::move(done_cb));
+    ctx->rendezvous()->RecvAsync(in_loop_parsed, args,
+                                 make_recv_callback(ctx, std::move(done)));
   }
 }