1 files changed, 59 insertions, 56 deletions

diff --git a/extensions/renderer/display_source_custom_bindings.cc b/extensions/renderer/display_source_custom_bindings.cc
index ab68553..1a021c1 100644
--- a/extensions/renderer/display_source_custom_bindings.cc
+++ b/extensions/renderer/display_source_custom_bindings.cc
@@ -67,13 +67,17 @@ v8::Local<v8::Value> GetChildValue(v8::Local<v8::Object> value,
   return v8::Null(isolate);
 }
 
+int32_t GetCallbackId() {
+  static int32_t sCallId = 0;
+  return ++sCallId;
+}
+
 }  // namespace
 
 void DisplaySourceCustomBindings::StartSession(
     const v8::FunctionCallbackInfo<v8::Value>& args) {
-  CHECK_EQ(2, args.Length());
+  CHECK_EQ(1, args.Length());
   CHECK(args[0]->IsObject());
-  CHECK(args[1]->IsFunction());
 
   v8::Isolate* isolate = context()->isolate();
   v8::Local<v8::Object> start_info = args[0].As<v8::Object>();
@@ -154,32 +158,28 @@ void DisplaySourceCustomBindings::StartSession(
     return;
   }
 
-  auto on_started_callback = base::Bind(
-      &DisplaySourceCustomBindings::OnSessionStarted,
-      weak_factory_.GetWeakPtr());
-  auto on_terminated_callback = base::Bind(
-      &DisplaySourceCustomBindings::OnSessionTerminated,
-      weak_factory_.GetWeakPtr());
-  auto on_error_callback = base::Bind(
-      &DisplaySourceCustomBindings::OnSessionError,
-      weak_factory_.GetWeakPtr());
-  session->SetCallbacks(on_started_callback,
-                        on_terminated_callback,
-                        on_error_callback);
-
-  CallbackInfo cb_info = GetCallbackInfo(kStarted, sink_id);
-  args.GetReturnValue().Set(static_cast<int32_t>(cb_info.call_id));
-  callbacks_.push_back(cb_info);
-
-  session->Start();
+  auto on_terminated_callback =
+      base::Bind(&DisplaySourceCustomBindings::OnSessionTerminated,
+                 weak_factory_.GetWeakPtr(), sink_id);
+  auto on_error_callback =
+      base::Bind(&DisplaySourceCustomBindings::OnSessionError,
+                 weak_factory_.GetWeakPtr(), sink_id);
+  session->SetNotificationCallbacks(on_terminated_callback, on_error_callback);
+
+  int32_t call_id = GetCallbackId();
+  args.GetReturnValue().Set(call_id);
+
+  auto on_call_completed =
+      base::Bind(&DisplaySourceCustomBindings::OnSessionStarted,
+                 weak_factory_.GetWeakPtr(), sink_id, call_id);
+  session->Start(on_call_completed);
   session_map_.insert(std::make_pair(sink_id, std::move(session)));
 }
 
 void DisplaySourceCustomBindings::TerminateSession(
     const v8::FunctionCallbackInfo<v8::Value>& args) {
-  CHECK_EQ(2, args.Length());
+  CHECK_EQ(1, args.Length());
   CHECK(args[0]->IsInt32());
-  CHECK(args[1]->IsFunction());
 
   v8::Isolate* isolate = context()->isolate();
   int sink_id = args[0]->ToInt32(args.GetIsolate())->Value();
@@ -190,39 +190,44 @@ void DisplaySourceCustomBindings::TerminateSession(
     return;
   }
 
-  if (session->state() == DisplaySourceSession::Terminating) {
+  DisplaySourceSession::State state = session->state();
+  DCHECK_NE(state, DisplaySourceSession::Idle);
+  if (state == DisplaySourceSession::Establishing) {
+    // 'session started' callback has not yet been invoked.
+    // This session is not existing for the user.
+    isolate->ThrowException(v8::Exception::Error(
+        v8::String::NewFromUtf8(isolate, kSessionNotFound)));
+    return;
+  }
+
+  if (state == DisplaySourceSession::Terminating) {
     isolate->ThrowException(v8::Exception::Error(v8::String::NewFromUtf8(
         isolate, kSessionAlreadyTerminating)));
     return;
   }
 
-  CallbackInfo cb_info = GetCallbackInfo(kTerminated, sink_id);
-  args.GetReturnValue().Set(static_cast<int32_t>(cb_info.call_id));
-  callbacks_.push_back(cb_info);
+  int32_t call_id = GetCallbackId();
+  args.GetReturnValue().Set(call_id);
 
+  auto on_call_completed =
+      base::Bind(&DisplaySourceCustomBindings::OnCallCompleted,
+                 weak_factory_.GetWeakPtr(), call_id);
   // The session will get removed from session_map_ in OnSessionTerminated.
-  session->Terminate();
+  session->Terminate(on_call_completed);
 }
 
-void DisplaySourceCustomBindings::CallCompletionCallback(
-    int sink_id,
-    CallbackType type,
+void DisplaySourceCustomBindings::OnCallCompleted(
+    int call_id,
+    bool success,
     const std::string& error_message) {
-  auto predicate = [sink_id, type](const CallbackInfo& info) -> bool {
-    return info.sink_id == sink_id && info.type == type;
-  };
-  auto it = std::find_if(callbacks_.begin(), callbacks_.end(), predicate);
-  if (it == callbacks_.end())
-    return;
-
   v8::Isolate* isolate = context()->isolate();
   ModuleSystem* module_system = context()->module_system();
   v8::HandleScope handle_scope(isolate);
   v8::Context::Scope context_scope(context()->v8_context());
 
   v8::Local<v8::Value> callback_args[2];
-  callback_args[0] = v8::Integer::New(isolate, it->call_id);
-  if (error_message.empty())
+  callback_args[0] = v8::Integer::New(isolate, call_id);
+  if (success)
     callback_args[1] = v8::Null(isolate);
   else
     callback_args[1] = v8::String::NewFromUtf8(isolate, error_message.c_str());
@@ -231,6 +236,19 @@ void DisplaySourceCustomBindings::CallCompletionCallback(
                                   callback_args);
 }
 
+void DisplaySourceCustomBindings::OnSessionStarted(
+    int sink_id,
+    int call_id,
+    bool success,
+    const std::string& error_message) {
+  CHECK(GetDisplaySession(sink_id));
+  if (!success) {
+    // Session has failed to start, removing it.
+    session_map_.erase(sink_id);
+  }
+  OnCallCompleted(call_id, success, error_message);
+}
+
 void DisplaySourceCustomBindings::DispatchSessionTerminated(int sink_id) const {
   v8::Isolate* isolate = context()->isolate();
   v8::HandleScope handle_scope(isolate);
@@ -272,32 +290,17 @@ DisplaySourceSession* DisplaySourceCustomBindings::GetDisplaySession(
   return nullptr;
 }
 
-void DisplaySourceCustomBindings::OnSessionStarted(int sink_id) {
-  CallCompletionCallback(sink_id, kStarted);
-}
-
 void DisplaySourceCustomBindings::OnSessionTerminated(int sink_id) {
-  DisplaySourceSession* session = GetDisplaySession(sink_id);
-  CHECK(session);
+  CHECK(GetDisplaySession(sink_id));
   session_map_.erase(sink_id);
   DispatchSessionTerminated(sink_id);
-  CallCompletionCallback(sink_id, kTerminated);
 }
 
 void DisplaySourceCustomBindings::OnSessionError(int sink_id,
                                                  DisplaySourceErrorType type,
                                                  const std::string& message) {
-  DisplaySourceSession* session = GetDisplaySession(sink_id);
-  CHECK(session);
+  CHECK(GetDisplaySession(sink_id));
   DispatchSessionError(sink_id, type, message);
 }
 
-DisplaySourceCustomBindings::CallbackInfo
-DisplaySourceCustomBindings::GetCallbackInfo(
-    DisplaySourceCustomBindings::CallbackType type,
-    int sink_id) const {
-  static int sCallId = 0;
-  return {type, sink_id, ++sCallId};
-}
-
 }  // extensions