Auto merge of #5884 - jgraham:webdriver_screenshot, r=jdm

This adds support for compositing to a PNG without actually quiting
the browser.

Cargo bits need to be updated after the upstream changes to rust-png land.

<!-- Reviewable:start -->
[<img src="https://reviewable.io/review_button.png" height=40 alt="Review on Reviewable"/>](https://reviewable.io/reviews/servo/servo/5884)
<!-- Reviewable:end -->

components/compositing/compositor.rs

@@ -92,6 +92,9 @@ pub struct IOCompositor<Window: WindowMethods> {
     /// A handle to the scrolling timer.
     scrolling_timer: ScrollingTimerProxy,
 
+    /// The type of composition to perform
+    composite_target: CompositeTarget,
+
     /// Tracks whether we should composite this frame.
     composition_request: CompositionRequest,
 
@@ -191,6 +194,38 @@ impl PipelineDetails {
     }
 }
 
+#[derive(Clone, Copy, PartialEq)]
+enum CompositeTarget {
+    /// Normal composition to a window
+    Window,
+
+    /// Compose as normal, but also return a PNG of the composed output
+    WindowAndPng,
+
+    /// Compose to a PNG, write it to disk, and then exit the browser (used for reftests)
+    PngFile
+}
+
+fn initialize_png(width: usize, height: usize) -> (Vec<gl::GLuint>, Vec<gl::GLuint>) {
+    let framebuffer_ids = gl::gen_framebuffers(1);
+    gl::bind_framebuffer(gl::FRAMEBUFFER, framebuffer_ids[0]);
+
+    let texture_ids = gl::gen_textures(1);
+    gl::bind_texture(gl::TEXTURE_2D, texture_ids[0]);
+
+    gl::tex_image_2d(gl::TEXTURE_2D, 0, gl::RGB as GLint, width as GLsizei,
+                     height as GLsizei, 0, gl::RGB, gl::UNSIGNED_BYTE, None);
+    gl::tex_parameter_i(gl::TEXTURE_2D, gl::TEXTURE_MAG_FILTER, gl::NEAREST as GLint);
+    gl::tex_parameter_i(gl::TEXTURE_2D, gl::TEXTURE_MIN_FILTER, gl::NEAREST as GLint);
+
+    gl::framebuffer_texture_2d(gl::FRAMEBUFFER, gl::COLOR_ATTACHMENT0, gl::TEXTURE_2D,
+                               texture_ids[0], 0);
+
+    gl::bind_texture(gl::TEXTURE_2D, 0);
+
+    (framebuffer_ids, texture_ids)
+}
+
 impl<Window: WindowMethods> IOCompositor<Window> {
     fn new(window: Rc<Window>,
            sender: Box<CompositorProxy+Send>,
@@ -205,6 +240,10 @@ impl<Window: WindowMethods> IOCompositor<Window> {
         // display list. This is only here because we don't have that logic in the painter yet.
         let window_size = window.framebuffer_size();
         let hidpi_factor = window.hidpi_factor();
+        let composite_target = match opts::get().output_file {
+            Some(_) => CompositeTarget::PngFile,
+            None => CompositeTarget::Window
+        };
         IOCompositor {
             window: window,
             port: receiver,
@@ -221,6 +260,7 @@ impl<Window: WindowMethods> IOCompositor<Window> {
             scrolling_timer: ScrollingTimerProxy::new(sender),
             composition_request: CompositionRequest::NoCompositingNecessary,
             pending_scroll_events: Vec::new(),
+            composite_target: composite_target,
             shutdown_state: ShutdownState::NotShuttingDown,
             page_zoom: ScaleFactor::new(1.0),
             viewport_zoom: ScaleFactor::new(1.0),
@@ -390,6 +430,11 @@ impl<Window: WindowMethods> IOCompositor<Window> {
                 self.window.set_cursor(cursor)
             }
 
+            (Msg::CreatePng(reply), ShutdownState::NotShuttingDown) => {
+                let img = self.composite_specific_target(CompositeTarget::WindowAndPng);
+                reply.send(img).unwrap();
+            }
+
             (Msg::PaintTaskExited(pipeline_id), ShutdownState::NotShuttingDown) => {
                 if self.pipeline_details.remove(&pipeline_id).is_none() {
                     panic!("Saw PaintTaskExited message from an unknown pipeline!");
@@ -414,7 +459,7 @@ impl<Window: WindowMethods> IOCompositor<Window> {
         self.window.set_ready_state(self.get_earliest_pipeline_ready_state());
 
         // If we're painting in headless mode, schedule a recomposite.
-        if opts::get().output_file.is_some() {
+        if let CompositeTarget::PngFile = self.composite_target {
             self.composite_if_necessary(CompositingReason::Headless)
         }
     }
@@ -1170,35 +1215,31 @@ impl<Window: WindowMethods> IOCompositor<Window> {
     }
 
     fn composite(&mut self) {
+        let target = self.composite_target;
+        self.composite_specific_target(target);
+    }
+
+    fn composite_specific_target(&mut self, target: CompositeTarget) -> Option<png::Image> {
         if !self.window.prepare_for_composite() {
-            return
+            return None
         }
 
-        let output_image = opts::get().output_file.is_some() &&
+        match target {
-                            self.is_ready_to_paint_image_output();
+            CompositeTarget::WindowAndPng | CompositeTarget::PngFile => {
+                if !self.is_ready_to_paint_image_output() {
+                    return None
+                }
+            },
+            _ => {}
+        }
 
-        let mut framebuffer_ids = vec!();
-        let mut texture_ids = vec!();
         let (width, height) =
             (self.window_size.width.get() as usize, self.window_size.height.get() as usize);
 
-        if output_image {
+        let (framebuffer_ids, texture_ids) = match target {
-            framebuffer_ids = gl::gen_framebuffers(1);
+            CompositeTarget::Window => (vec!(), vec!()),
-            gl::bind_framebuffer(gl::FRAMEBUFFER, framebuffer_ids[0]);
+            _ => initialize_png(width, height)
-
+        };
-            texture_ids = gl::gen_textures(1);
-            gl::bind_texture(gl::TEXTURE_2D, texture_ids[0]);
-
-            gl::tex_image_2d(gl::TEXTURE_2D, 0, gl::RGB as GLint, width as GLsizei,
-                             height as GLsizei, 0, gl::RGB, gl::UNSIGNED_BYTE, None);
-            gl::tex_parameter_i(gl::TEXTURE_2D, gl::TEXTURE_MAG_FILTER, gl::NEAREST as GLint);
-            gl::tex_parameter_i(gl::TEXTURE_2D, gl::TEXTURE_MIN_FILTER, gl::NEAREST as GLint);
-
-            gl::framebuffer_texture_2d(gl::FRAMEBUFFER, gl::COLOR_ATTACHMENT0, gl::TEXTURE_2D,
-                                       texture_ids[0], 0);
-
-            gl::bind_texture(gl::TEXTURE_2D, 0);
-        }
 
         profile(ProfilerCategory::Compositing, None, self.time_profiler_chan.clone(), || {
             debug!("compositor: compositing");
@@ -1219,41 +1260,24 @@ impl<Window: WindowMethods> IOCompositor<Window> {
             }
         });
 
-        if output_image {
+        let rv = match target {
-            let path = opts::get().output_file.as_ref().unwrap();
+            CompositeTarget::Window => None,
-            let mut pixels = gl::read_pixels(0, 0,
+            CompositeTarget::WindowAndPng => {
-                                             width as gl::GLsizei,
+                Some(self.draw_png(framebuffer_ids, texture_ids, width, height))
-                                             height as gl::GLsizei,
-                                             gl::RGB, gl::UNSIGNED_BYTE);
-
-            gl::bind_framebuffer(gl::FRAMEBUFFER, 0);
-
-            gl::delete_buffers(&texture_ids);
-            gl::delete_frame_buffers(&framebuffer_ids);
-
-            // flip image vertically (texture is upside down)
-            let orig_pixels = pixels.clone();
-            let stride = width * 3;
-            for y in 0..height {
-                let dst_start = y * stride;
-                let src_start = (height - y - 1) * stride;
-                let src_slice = &orig_pixels[src_start .. src_start + stride];
-                copy_memory(&src_slice[..stride],
-                            &mut pixels[dst_start .. dst_start + stride]);
             }
-            let mut img = png::Image {
+            CompositeTarget::PngFile => {
-                width: width as u32,
+                let mut img = self.draw_png(framebuffer_ids, texture_ids, width, height);
-                height: height as u32,
+                let path = opts::get().output_file.as_ref().unwrap();
-                pixels: png::PixelsByColorType::RGB8(pixels),
+                let res = png::store_png(&mut img, &path);
-            };
+                assert!(res.is_ok());
-            let res = png::store_png(&mut img, &path);
-            assert!(res.is_ok());
 
-            debug!("shutting down the constellation after generating an output file");
+                debug!("shutting down the constellation after generating an output file");
-            let ConstellationChan(ref chan) = self.constellation_chan;
+                let ConstellationChan(ref chan) = self.constellation_chan;
-            chan.send(ConstellationMsg::Exit).unwrap();
+                chan.send(ConstellationMsg::Exit).unwrap();
-            self.shutdown_state = ShutdownState::ShuttingDown;
+                self.shutdown_state = ShutdownState::ShuttingDown;
-        }
+                None
+            }
+        };
 
         // Perform the page flip. This will likely block for a while.
         self.window.present();
@@ -1263,6 +1287,35 @@ impl<Window: WindowMethods> IOCompositor<Window> {
         self.composition_request = CompositionRequest::NoCompositingNecessary;
         self.process_pending_scroll_events();
         self.process_animations();
+        rv
+    }
+
+    fn draw_png(&self, framebuffer_ids: Vec<gl::GLuint>, texture_ids: Vec<gl::GLuint>, width: usize, height: usize) -> png::Image {
+        let mut pixels = gl::read_pixels(0, 0,
+                                         width as gl::GLsizei,
+                                         height as gl::GLsizei,
+                                         gl::RGB, gl::UNSIGNED_BYTE);
+
+        gl::bind_framebuffer(gl::FRAMEBUFFER, 0);
+
+        gl::delete_buffers(&texture_ids);
+        gl::delete_frame_buffers(&framebuffer_ids);
+
+        // flip image vertically (texture is upside down)
+        let orig_pixels = pixels.clone();
+        let stride = width * 3;
+        for y in 0..height {
+            let dst_start = y * stride;
+            let src_start = (height - y - 1) * stride;
+            let src_slice = &orig_pixels[src_start .. src_start + stride];
+            copy_memory(&src_slice[..stride],
+                        &mut pixels[dst_start .. dst_start + stride]);
+        }
+        png::Image {
+            width: width as u32,
+            height: height as u32,
+            pixels: png::PixelsByColorType::RGB8(pixels),
+        }
     }
 
     fn composite_if_necessary(&mut self, reason: CompositingReason) {

components/compositing/compositor_task.rs

@@ -23,6 +23,7 @@ use msg::constellation_msg::{AnimationState, ConstellationChan, PipelineId};
 use msg::constellation_msg::{Key, KeyState, KeyModifiers};
 use profile_traits::mem;
 use profile_traits::time;
+use png;
 use std::sync::mpsc::{channel, Sender, Receiver};
 use std::fmt::{Error, Formatter, Debug};
 use std::rc::Rc;
@@ -218,6 +219,8 @@ pub enum Msg {
     KeyEvent(Key, KeyState, KeyModifiers),
     /// Changes the cursor.
     SetCursor(Cursor),
+    /// Composite to a PNG file and return the Image over a passed channel.
+    CreatePng(Sender<Option<png::Image>>),
     /// Informs the compositor that the paint task for the given pipeline has exited.
     PaintTaskExited(PipelineId),
     /// Alerts the compositor that the viewport has been constrained in some manner
@@ -247,6 +250,7 @@ impl Debug for Msg {
             Msg::RecompositeAfterScroll => write!(f, "RecompositeAfterScroll"),
             Msg::KeyEvent(..) => write!(f, "KeyEvent"),
             Msg::SetCursor(..) => write!(f, "SetCursor"),
+            Msg::CreatePng(..) => write!(f, "CreatePng"),
             Msg::PaintTaskExited(..) => write!(f, "PaintTaskExited"),
             Msg::ViewportConstrained(..) => write!(f, "ViewportConstrained"),
         }

components/compositing/constellation.rs

@@ -413,6 +413,9 @@ impl<LTF: LayoutTaskFactory, STF: ScriptTaskFactory> Constellation<LTF, STF> {
                 };
                 sender.send(result).unwrap();
             }
+            ConstellationMsg::CompositePng(reply) => {
+                self.compositor_proxy.send(CompositorMsg::CreatePng(reply));
+            }
             ConstellationMsg::WebDriverCommand(pipeline_id,
                                                command) => {
                 debug!("constellation got webdriver command message");

components/compositing/headless.rs

@@ -108,8 +108,9 @@ impl CompositorEventListener for NullCompositor {
             Msg::ChangePageUrl(..) |
             Msg::KeyEvent(..) |
             Msg::SetCursor(..) |
-            Msg::PaintTaskExited(..) |
             Msg::ViewportConstrained(..) => {}
+            Msg::CreatePng(..) |
+            Msg::PaintTaskExited(..) => {}
         }
         true
     }

components/msg/Cargo.toml

@@ -32,6 +32,9 @@ git = "https://github.com/servo/rust-core-foundation"
 [dependencies.io_surface]
 git = "https://github.com/servo/rust-io-surface"
 
+[dependencies.png]
+git = "https://github.com/servo/rust-png"
+
 [dependencies]
 url = "0.2.16"
 bitflags = "*"

components/msg/constellation_msg.rs

@@ -11,6 +11,7 @@ use geom::scale_factor::ScaleFactor;
 use hyper::header::Headers;
 use hyper::method::Method;
 use layers::geometry::DevicePixel;
+use png;
 use util::cursor::Cursor;
 use util::geometry::{PagePx, ViewportPx};
 use std::sync::mpsc::{channel, Sender, Receiver};
@@ -233,10 +234,12 @@ pub enum Msg {
     Focus(PipelineId),
     /// Requests that the constellation retrieve the current contents of the clipboard
     GetClipboardContents(Sender<String>),
-    // Dispatch a webdriver command
+    /// Dispatch a webdriver command
     WebDriverCommand(PipelineId, WebDriverScriptCommand),
     /// Notifies the constellation that the viewport has been constrained in some manner
     ViewportConstrained(PipelineId, ViewportConstraints),
+    /// Create a PNG of the window contents
+    CompositePng(Sender<Option<png::Image>>)
 }
 
 #[derive(Clone, Eq, PartialEq)]

components/msg/lib.rs

@@ -7,6 +7,7 @@ extern crate azure;
 extern crate geom;
 extern crate hyper;
 extern crate layers;
+extern crate png;
 extern crate util;
 extern crate url;
 extern crate style;

components/script/script_task.rs

@@ -724,9 +724,8 @@ impl ScriptTask {
                 self.handle_update_subpage_id(containing_pipeline_id, old_subpage_id, new_subpage_id),
             ConstellationControlMsg::FocusIFrame(containing_pipeline_id, subpage_id) =>
                 self.handle_focus_iframe_msg(containing_pipeline_id, subpage_id),
-            ConstellationControlMsg::WebDriverCommand(pipeline_id, msg) => {
+            ConstellationControlMsg::WebDriverCommand(pipeline_id, msg) =>
-                self.handle_webdriver_msg(pipeline_id, msg);
+                self.handle_webdriver_msg(pipeline_id, msg),
-            }
             ConstellationControlMsg::TickAllAnimations(pipeline_id) =>
                 self.handle_tick_all_animations(pipeline_id),
         }

components/servo/Cargo.lock

@@ -751,6 +751,7 @@ dependencies = [
  "hyper 0.3.16 (registry+https://github.com/rust-lang/crates.io-index)",
  "io_surface 0.1.0 (git+https://github.com/servo/rust-io-surface)",
  "layers 0.1.0 (git+https://github.com/servo/rust-layers)",
+ "png 0.1.0 (git+https://github.com/servo/rust-png)",
  "style 0.0.1",
  "url 0.2.31 (registry+https://github.com/rust-lang/crates.io-index)",
  "util 0.0.1",
@@ -1251,7 +1252,7 @@ dependencies = [
 [[package]]
 name = "webdriver"
 version = "0.1.0"
-source = "git+https://github.com/jgraham/webdriver-rust.git#66547888f47bae7e938a92af4586276479343216"
+source = "git+https://github.com/jgraham/webdriver-rust.git#c2038b4195ee8cd982079cc48d6a9d039f59f1fb"
 dependencies = [
  "hyper 0.3.16 (registry+https://github.com/rust-lang/crates.io-index)",
  "log 0.3.1 (registry+https://github.com/rust-lang/crates.io-index)",
@@ -1265,6 +1266,7 @@ name = "webdriver_server"
 version = "0.0.1"
 dependencies = [
  "msg 0.0.1",
+ "png 0.1.0 (git+https://github.com/servo/rust-png)",
  "rustc-serialize 0.3.13 (registry+https://github.com/rust-lang/crates.io-index)",
  "url 0.2.31 (registry+https://github.com/rust-lang/crates.io-index)",
  "util 0.0.1",

components/webdriver_server/Cargo.toml

@@ -19,7 +19,10 @@ path = "../webdriver_traits"
 [dependencies.webdriver]
 git = "https://github.com/jgraham/webdriver-rust.git"
 
+[dependencies.png]
+git = "https://github.com/servo/rust-png"
+
 [dependencies]
 rustc-serialize = "0.3.4"
 url = "0.2.16"
-uuid = "*"
+uuid = "*"

components/webdriver_server/lib.rs

@@ -12,6 +12,7 @@ extern crate log;
 
 extern crate webdriver;
 extern crate msg;
+extern crate png;
 extern crate url;
 extern crate util;
 extern crate rustc_serialize;
@@ -35,9 +36,12 @@ use uuid::Uuid;
 
 use std::borrow::ToOwned;
 use rustc_serialize::json::{Json, ToJson};
+use rustc_serialize::base64::{Config, ToBase64, CharacterSet, Newline};
 use std::collections::BTreeMap;
 use std::net::SocketAddr;
 
+use std::thread::sleep_ms;
+
 pub fn start_server(port: u16, constellation_chan: ConstellationChan) {
     let handler = Handler::new(constellation_chan);
 
@@ -171,6 +175,46 @@ impl Handler {
                                               "Unsupported return type"))
         }
     }
+
+
+    fn handle_take_screenshot(&self) -> WebDriverResult<WebDriverResponse> {
+        let mut img = None;
+
+        let interval = 20;
+        let iterations = 30_000 / interval;
+
+        for _ in 0..iterations {
+            let (sender, reciever) = channel();
+            let ConstellationChan(ref const_chan) = self.constellation_chan;
+            const_chan.send(ConstellationMsg::CompositePng(sender)).unwrap();
+
+            if let Some(x) = reciever.recv().unwrap() {
+                img = Some(x);
+                break;
+            };
+
+            sleep_ms(interval)
+        }
+
+        if img.is_none() {
+            return Err(WebDriverError::new(ErrorStatus::Timeout,
+                                           "Taking screenshot timed out"));
+        }
+
+        let img_vec = match png::to_vec(&mut img.unwrap()) {
+           Ok(x) => x,
+           Err(_) => return Err(WebDriverError::new(ErrorStatus::UnknownError,
+                                                    "Taking screenshot failed"))
+        };
+        let config = Config {
+            char_set:CharacterSet::Standard,
+            newline: Newline::LF,
+            pad: true,
+            line_length: None
+        };
+        let encoded = img_vec.to_base64(config);
+        Ok(WebDriverResponse::Generic(ValueResponse::new(encoded.to_json())))
+    }
 }
 
 impl WebDriverHandler for Handler {
@@ -185,6 +229,7 @@ impl WebDriverHandler for Handler {
             WebDriverCommand::GetWindowHandle => self.handle_get_window_handle(),
             WebDriverCommand::GetWindowHandles => self.handle_get_window_handles(),
             WebDriverCommand::ExecuteScript(ref x) => self.handle_execute_script(x),
+            WebDriverCommand::TakeScreenshot => self.handle_take_screenshot(),
             _ => Err(WebDriverError::new(ErrorStatus::UnsupportedOperation,
                                          "Command not implemented"))
         }

ports/cef/Cargo.lock

@@ -752,6 +752,7 @@ dependencies = [
  "hyper 0.3.16 (registry+https://github.com/rust-lang/crates.io-index)",
  "io_surface 0.1.0 (git+https://github.com/servo/rust-io-surface)",
  "layers 0.1.0 (git+https://github.com/servo/rust-layers)",
+ "png 0.1.0 (git+https://github.com/servo/rust-png)",
  "style 0.0.1",
  "url 0.2.31 (registry+https://github.com/rust-lang/crates.io-index)",
  "util 0.0.1",
@@ -1235,7 +1236,7 @@ dependencies = [
 [[package]]
 name = "webdriver"
 version = "0.1.0"
-source = "git+https://github.com/jgraham/webdriver-rust.git#66547888f47bae7e938a92af4586276479343216"
+source = "git+https://github.com/jgraham/webdriver-rust.git#c2038b4195ee8cd982079cc48d6a9d039f59f1fb"
 dependencies = [
  "hyper 0.3.16 (registry+https://github.com/rust-lang/crates.io-index)",
  "log 0.3.1 (registry+https://github.com/rust-lang/crates.io-index)",
@@ -1249,6 +1250,7 @@ name = "webdriver_server"
 version = "0.0.1"
 dependencies = [
  "msg 0.0.1",
+ "png 0.1.0 (git+https://github.com/servo/rust-png)",
  "rustc-serialize 0.3.13 (registry+https://github.com/rust-lang/crates.io-index)",
  "url 0.2.31 (registry+https://github.com/rust-lang/crates.io-index)",
  "util 0.0.1",

ports/gonk/Cargo.lock

@@ -725,6 +725,7 @@ dependencies = [
  "hyper 0.3.16 (registry+https://github.com/rust-lang/crates.io-index)",
  "io_surface 0.1.0 (git+https://github.com/servo/rust-io-surface)",
  "layers 0.1.0 (git+https://github.com/servo/rust-layers)",
+ "png 0.1.0 (git+https://github.com/servo/rust-png)",
  "style 0.0.1",
  "url 0.2.31 (registry+https://github.com/rust-lang/crates.io-index)",
  "util 0.0.1",
@@ -1207,7 +1208,7 @@ dependencies = [
 [[package]]
 name = "webdriver"
 version = "0.1.0"
-source = "git+https://github.com/jgraham/webdriver-rust.git#66547888f47bae7e938a92af4586276479343216"
+source = "git+https://github.com/jgraham/webdriver-rust.git#c2038b4195ee8cd982079cc48d6a9d039f59f1fb"
 dependencies = [
  "hyper 0.3.16 (registry+https://github.com/rust-lang/crates.io-index)",
  "log 0.3.1 (registry+https://github.com/rust-lang/crates.io-index)",
@@ -1221,6 +1222,7 @@ name = "webdriver_server"
 version = "0.0.1"
 dependencies = [
  "msg 0.0.1",
+ "png 0.1.0 (git+https://github.com/servo/rust-png)",
  "rustc-serialize 0.3.13 (registry+https://github.com/rust-lang/crates.io-index)",
  "url 0.2.31 (registry+https://github.com/rust-lang/crates.io-index)",
  "util 0.0.1",