1// Copyright 2014 The Chromium Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5#include "remoting/host/video_frame_recorder.h"
6
7#include "base/bind.h"
8#include "base/location.h"
9#include "base/single_thread_task_runner.h"
10#include "base/stl_util.h"
11#include "base/thread_task_runner_handle.h"
12#include "remoting/codec/video_encoder.h"
13#include "remoting/proto/video.pb.h"
14#include "third_party/webrtc/modules/desktop_capture/desktop_frame.h"
15#include "third_party/webrtc/modules/desktop_capture/desktop_geometry.h"
16#include "third_party/webrtc/modules/desktop_capture/desktop_region.h"
17
18namespace remoting {
19
20static int64_t FrameContentSize(const webrtc::DesktopFrame* frame) {
21  DCHECK_GT(frame->stride(), 0);
22  return frame->stride() * frame->size().height();
23}
24
25// VideoEncoder wrapper used to intercept frames passed to a real VideoEncoder.
26class VideoFrameRecorder::RecordingVideoEncoder : public VideoEncoder {
27 public:
28  RecordingVideoEncoder(scoped_ptr<VideoEncoder> encoder,
29                        scoped_refptr<base::TaskRunner> recorder_task_runner,
30                        base::WeakPtr<VideoFrameRecorder> recorder);
31
32  base::WeakPtr<RecordingVideoEncoder> AsWeakPtr();
33
34  void set_enable_recording(bool enable_recording) {
35    DCHECK(!encoder_task_runner_.get() ||
36           encoder_task_runner_->BelongsToCurrentThread());
37    enable_recording_ = enable_recording;
38  }
39
40  // remoting::VideoEncoder interface.
41  virtual void SetLosslessEncode(bool want_lossless) OVERRIDE;
42  virtual void SetLosslessColor(bool want_lossless) OVERRIDE;
43  virtual scoped_ptr<VideoPacket> Encode(
44      const webrtc::DesktopFrame& frame) OVERRIDE;
45
46 private:
47  scoped_ptr<VideoEncoder> encoder_;
48  scoped_refptr<base::TaskRunner> recorder_task_runner_;
49  base::WeakPtr<VideoFrameRecorder> recorder_;
50
51  bool enable_recording_;
52  scoped_refptr<base::SingleThreadTaskRunner> encoder_task_runner_;
53
54  base::WeakPtrFactory<RecordingVideoEncoder> weak_factory_;
55
56  DISALLOW_COPY_AND_ASSIGN(RecordingVideoEncoder);
57};
58
59VideoFrameRecorder::RecordingVideoEncoder::RecordingVideoEncoder(
60    scoped_ptr<VideoEncoder> encoder,
61    scoped_refptr<base::TaskRunner> recorder_task_runner,
62    base::WeakPtr<VideoFrameRecorder> recorder)
63    : encoder_(encoder.Pass()),
64      recorder_task_runner_(recorder_task_runner),
65      recorder_(recorder),
66      enable_recording_(false),
67      weak_factory_(this) {
68  DCHECK(encoder_);
69  DCHECK(recorder_task_runner_.get());
70}
71
72base::WeakPtr<VideoFrameRecorder::RecordingVideoEncoder>
73VideoFrameRecorder::RecordingVideoEncoder::AsWeakPtr() {
74  return weak_factory_.GetWeakPtr();
75}
76
77void VideoFrameRecorder::RecordingVideoEncoder::SetLosslessEncode(
78    bool want_lossless) {
79  encoder_->SetLosslessEncode(want_lossless);
80}
81
82void VideoFrameRecorder::RecordingVideoEncoder::SetLosslessColor(
83    bool want_lossless) {
84  encoder_->SetLosslessColor(want_lossless);
85}
86
87scoped_ptr<VideoPacket> VideoFrameRecorder::RecordingVideoEncoder::Encode(
88    const webrtc::DesktopFrame& frame) {
89  // If this is the first Encode() then store the TaskRunner and inform the
90  // VideoFrameRecorder so it can post set_enable_recording() on it.
91  if (!encoder_task_runner_.get()) {
92    encoder_task_runner_ = base::ThreadTaskRunnerHandle::Get();
93    recorder_task_runner_->PostTask(FROM_HERE,
94        base::Bind(&VideoFrameRecorder::SetEncoderTaskRunner,
95                   recorder_,
96                   encoder_task_runner_));
97  }
98
99  DCHECK(encoder_task_runner_->BelongsToCurrentThread());
100
101  if (enable_recording_) {
102    // Copy the frame and post it to the VideoFrameRecorder to store.
103    scoped_ptr<webrtc::DesktopFrame> frame_copy(
104        new webrtc::BasicDesktopFrame(frame.size()));
105    *frame_copy->mutable_updated_region() = frame.updated_region();
106    frame_copy->set_dpi(frame.dpi());
107    frame_copy->CopyPixelsFrom(frame.data(),
108                               frame.stride(),
109                               webrtc::DesktopRect::MakeSize(frame.size()));
110    recorder_task_runner_->PostTask(FROM_HERE,
111        base::Bind(&VideoFrameRecorder::RecordFrame,
112                   recorder_,
113                   base::Passed(&frame_copy)));
114  }
115
116  return encoder_->Encode(frame);
117}
118
119VideoFrameRecorder::VideoFrameRecorder()
120    : content_bytes_(0),
121      max_content_bytes_(0),
122      enable_recording_(false),
123      weak_factory_(this) {
124}
125
126VideoFrameRecorder::~VideoFrameRecorder() {
127  DetachVideoEncoderWrapper();
128}
129
130scoped_ptr<VideoEncoder> VideoFrameRecorder::WrapVideoEncoder(
131    scoped_ptr<VideoEncoder> encoder) {
132  DCHECK(!encoder_task_runner_.get());
133  DCHECK(!caller_task_runner_.get());
134  caller_task_runner_ = base::ThreadTaskRunnerHandle::Get();
135
136  scoped_ptr<RecordingVideoEncoder> recording_encoder(
137      new RecordingVideoEncoder(encoder.Pass(),
138                                caller_task_runner_,
139                                weak_factory_.GetWeakPtr()));
140  recording_encoder_ = recording_encoder->AsWeakPtr();
141
142  return recording_encoder.PassAs<VideoEncoder>();
143}
144
145void VideoFrameRecorder::DetachVideoEncoderWrapper() {
146  DCHECK(!caller_task_runner_.get() ||
147         caller_task_runner_->BelongsToCurrentThread());
148
149  // Immediately detach the wrapper from this recorder.
150  weak_factory_.InvalidateWeakPtrs();
151
152  // Clean up any pending recorded frames.
153  STLDeleteElements(&recorded_frames_);
154  content_bytes_ = 0;
155
156  // Tell the wrapper to stop recording and posting frames to us.
157  if (encoder_task_runner_.get()) {
158    encoder_task_runner_->PostTask(FROM_HERE,
159        base::Bind(&RecordingVideoEncoder::set_enable_recording,
160                   recording_encoder_, false));
161  }
162
163  // Detach this recorder from the calling and encode threads.
164  caller_task_runner_ = NULL;
165  encoder_task_runner_ = NULL;
166}
167
168void VideoFrameRecorder::SetEnableRecording(bool enable_recording) {
169  DCHECK(!caller_task_runner_.get() ||
170         caller_task_runner_->BelongsToCurrentThread());
171
172  if (enable_recording_ == enable_recording) {
173    return;
174  }
175  enable_recording_ = enable_recording;
176
177  if (encoder_task_runner_.get()) {
178    encoder_task_runner_->PostTask(FROM_HERE,
179        base::Bind(&RecordingVideoEncoder::set_enable_recording,
180                   recording_encoder_,
181                   enable_recording_));
182  }
183}
184
185void VideoFrameRecorder::SetMaxContentBytes(int64_t max_content_bytes) {
186  DCHECK(!caller_task_runner_.get() ||
187         caller_task_runner_->BelongsToCurrentThread());
188  DCHECK_GE(max_content_bytes, 0);
189
190  max_content_bytes_ = max_content_bytes;
191}
192
193scoped_ptr<webrtc::DesktopFrame> VideoFrameRecorder::NextFrame() {
194  DCHECK(caller_task_runner_->BelongsToCurrentThread());
195
196  scoped_ptr<webrtc::DesktopFrame> frame;
197  if (!recorded_frames_.empty()) {
198    frame.reset(recorded_frames_.front());
199    recorded_frames_.pop_front();
200    content_bytes_ -= FrameContentSize(frame.get());
201    DCHECK_GE(content_bytes_, 0);
202  }
203
204  return frame.Pass();
205}
206
207void VideoFrameRecorder::SetEncoderTaskRunner(
208    scoped_refptr<base::TaskRunner> task_runner) {
209  DCHECK(caller_task_runner_->BelongsToCurrentThread());
210  DCHECK(!encoder_task_runner_.get());
211  DCHECK(task_runner.get());
212
213  encoder_task_runner_ = task_runner;
214
215  // If the caller already enabled recording, inform the recording encoder.
216  if (enable_recording_ && encoder_task_runner_.get()) {
217    encoder_task_runner_->PostTask(FROM_HERE,
218        base::Bind(&RecordingVideoEncoder::set_enable_recording,
219                   recording_encoder_,
220                   enable_recording_));
221  }
222}
223
224void VideoFrameRecorder::RecordFrame(scoped_ptr<webrtc::DesktopFrame> frame) {
225  DCHECK(caller_task_runner_->BelongsToCurrentThread());
226
227  int64_t frame_bytes = FrameContentSize(frame.get());
228  DCHECK_GE(frame_bytes, 0);
229
230  // Purge existing frames until there is space for the new one.
231  while (content_bytes_ + frame_bytes > max_content_bytes_ &&
232         !recorded_frames_.empty()) {
233    scoped_ptr<webrtc::DesktopFrame> drop_frame(recorded_frames_.front());
234    recorded_frames_.pop_front();
235    content_bytes_ -= FrameContentSize(drop_frame.get());
236    DCHECK_GE(content_bytes_, 0);
237  }
238
239  // If the frame is still too big, ignore it.
240  if (content_bytes_ + frame_bytes > max_content_bytes_) {
241    return;
242  }
243
244  // Store the frame and update the content byte count.
245  recorded_frames_.push_back(frame.release());
246  content_bytes_ += frame_bytes;
247}
248
249}  // namespace remoting
250