1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
//! A thread pool that can process file requests and send data to the socket
//! with zero copy (using sendfile).
//!
//! Use `DiskPool` structure to request file operations.
//!
//! # Example
//!
//! ```rust,ignore
//!     let pool = DiskPool::new(CpuPool::new(40));
//!     pool.send("file", socket)
//! ```
//!
//! # Settings
//!
//! It's recommended to make large number of threads in the pool
//! for three reasons:
//!
//! 1. To make use of device parallelism
//! 2. To allow kernel to merge some disk operations
//! 3. To fix head of line blocking when some request reach disk but others
//!    could be served immediately from cache (we don't know which ones are
//!    cached, so we run all of them in a pool)
#![warn(missing_docs)]

extern crate libc;
extern crate futures;
extern crate tokio_io;
extern crate tokio_core;
extern crate futures_cpupool;

use std::io;
use std::mem;
use std::cmp::min;
use std::fs::File;
use std::path::PathBuf;

#[cfg(windows)]
use std::sync::Mutex;

use futures::{Future, Poll, Async, BoxFuture, finished, failed};
use futures_cpupool::{CpuPool, CpuFuture};
use tokio_io::AsyncWrite;
use tokio_core::net::TcpStream;

/// A reference to a thread pool for disk operations
#[derive(Clone)]
pub struct DiskPool {
    pool: CpuPool,
}

/// This trait represents anything that can open the file
///
/// You can convert anything that is `AsRef<Path>` to this trait and
/// it will just open a file at specified path.
/// But often you want some checks of permissions or visibility of the file
/// while opening it. You can't do even `stat()` or `open()` in main loop
/// because even such a simple operation can potentially block for indefinite
/// period of time.
///
/// So file opener could be anything that validates a path,
/// caches file descriptor, and in the result returns a file.
pub trait FileOpener: Send + 'static {
    /// Read file from cache
    ///
    /// Note: this can be both positive and negative cache
    ///
    /// You don't have to implement this method if you don't have in-memory
    /// cache of files
    fn from_cache(&mut self) -> Option<Result<&[u8], io::Error>> {
        None
    }
    /// Open the file
    ///
    /// This function is called in disk thread
    fn open(&mut self) -> Result<(&FileReader, u64), io::Error>;
}


/// This trait represents file that can atomically be read at specific point
///
/// For unix we implement it for `File + AsRawFd` with `pread()` system call.
/// For windows we're implementing it for `Mutex<File>` and use seek.
pub trait FileReader {
    /// Read file at specified location
    ///
    /// This corresponds to the `pread` system call on most unix systems
    fn read_at(&self, offset: u64, buf: &mut [u8]) -> io::Result<usize>;
}

/// Trait that represents something that can be converted into a file
/// FileOpener
///
/// This is very similar to IntoIterator or IntoFuture and used in similar
/// way.
///
/// Note unlike methods in FileOpener itself this trait is executed in
/// caller thread, **not** in disk thread.
pub trait IntoFileOpener: Send {
    /// The final type returned after conversion
    type Opener: FileOpener + Send + 'static;
    /// Convert the type into a file opener
    fn into_file_opener(self) -> Self::Opener;
}

/// File opener implementation that opens specified file path directly
#[derive(Debug)]
#[cfg(unix)]
pub struct PathOpener(PathBuf, Option<(File, u64)>);

/// File opener implementation that opens specified file path directly
#[derive(Debug)]
#[cfg(windows)]
pub struct PathOpener(PathBuf, Option<(Mutex<File>, u64)>);

/// A trait that represents anything that file can be sent to
///
/// The trait is implemented for TcpStream right away but you might want
/// to implement your own thing, for example to prepend the data with file
/// length header
pub trait Destination: AsyncWrite + Send {

    /// This method does the actual sendfile call
    ///
    /// Note: this method is called in other thread
    fn write_file<O: FileOpener>(&mut self, file: &mut Sendfile<O>)
        -> Result<usize, io::Error>;

    /// Test whether this socket is ready to be written to or not.
    ///
    /// If socket isn't writable current taks must be scheduled to get a
    /// notification when socket does become writable.
    fn poll_write(&self) -> Async<()>;
}

/// A structure that tracks progress of sending a file
pub struct Sendfile<O: FileOpener + Send + 'static> {
    file: O,
    pool: DiskPool,
    cached: bool,
    offset: u64,
    size: u64,
}

// Todo return non-boxed future
// /// Future that is returned from `DiskPool::send`
// type SendfileFuture<D> = futures_cpupool::CpuFuture<D, io::Error>;

/// Future returned by `Sendfile::write_into()`
pub struct WriteFile<F: FileOpener, D: Destination>(DiskPool, WriteState<F, D>)
    where F: Send + 'static, D: Send + 'static;

enum WriteState<F: FileOpener, D: Destination> {
    Mem(Sendfile<F>, D),
    WaitSend(CpuFuture<(Sendfile<F>, D), io::Error>),
    WaitWrite(Sendfile<F>, D),
    Empty,
}

impl<T: Into<PathBuf> + Send> IntoFileOpener for T {
    type Opener = PathOpener;
    fn into_file_opener(self) -> PathOpener {
        PathOpener(self.into(), None)
    }
}

#[cfg(unix)]
impl FileOpener for PathOpener {
    fn open(&mut self) -> Result<(&FileReader, u64), io::Error> {
        if self.1.is_none() {
            let file = File::open(&self.0)?;
            let meta = file.metadata()?;
            if !meta.file_type().is_file() {
                return Err(io::Error::new(io::ErrorKind::Other,
                    "Not a regular file"));
            }
            self.1 = Some((file, meta.len()));
        }
        Ok(self.1.as_ref().map(|&(ref f, s)| (f as &FileReader, s)).unwrap())
    }
}

#[cfg(windows)]
impl FileOpener for PathOpener {
    fn open(&mut self) -> Result<(&FileReader, u64), io::Error> {
        if self.1.is_none() {
            let file = File::open(&self.0)?;
            let meta = file.metadata()?;
            if !meta.file_type().is_file() {
                return Err(io::Error::new(io::ErrorKind::Other,
                    "Not a regular file"));
            }
            self.1 = Some((Mutex::new(file), meta.len()));
        }
        Ok(self.1.as_ref().map(|&(ref f, s)| (f as &FileReader, s)).unwrap())
    }
}

impl DiskPool {
    /// Create a disk pool that sends its tasks into the CpuPool
    pub fn new(pool: CpuPool) -> DiskPool {
        DiskPool {
            pool: pool,
        }
    }
    /// Start a file send operation
    pub fn open<F>(&self, file: F)
        // TODO(tailhook) unbox a future
        -> BoxFuture<Sendfile<F::Opener>, io::Error>
        where F: IntoFileOpener + Send + Sized + 'static,
    {
        let mut file = file.into_file_opener();
        let cached_size = match file.from_cache() {
            Some(Ok(cache_ref)) => {
                Some(cache_ref.len() as u64)
            }
            Some(Err(e)) => {
                return failed(e).boxed();
            }
            None => None,
        };
        let pool = self.clone();
        if let Some(size) = cached_size {
            finished(Sendfile {
                file: file,
                pool: pool,
                cached: true,
                offset: 0,
                size: size,
            }).boxed()
        } else {
            self.pool.spawn_fn(move || {
                let (_, size) = file.open()?;
                let file = Sendfile {
                    file: file,
                    pool: pool,
                    cached: false,
                    offset: 0,
                    size: size,
                };
                Ok(file)
            }).boxed()
        }
    }
    /// A shortcut method to send whole file without headers
    pub fn send<F, D>(&self, file: F, destination: D)
        -> futures::BoxFuture<D, io::Error>
        where F: IntoFileOpener + Send + Sized + 'static,
              D: Destination + Send + Sized + 'static,
    {
        self.open(file).and_then(|file| file.write_into(destination)).boxed()
    }
}

impl Destination for TcpStream {
    fn write_file<O: FileOpener>(&mut self, file: &mut Sendfile<O>)
        -> Result<usize, io::Error>
    {
        let (file_ref, size) = file.file.open()?;
        let mut buf = [0u8; 65536];
        let max_bytes = min(size.saturating_sub(file.offset), 65536) as usize;
        let nbytes = file_ref.read_at(file.offset, &mut buf[..max_bytes])?;
        if nbytes == 0 {
            return Err(io::ErrorKind::UnexpectedEof.into())
        }
        io::Write::write(self, &buf[..nbytes])
    }
    fn poll_write(&self) -> Async<()> {
        <TcpStream>::poll_write(self)
    }
}

impl<O: FileOpener> Sendfile<O> {
    /// Returns full size of the file
    ///
    /// Note that if file changes while we are reading it, we may not be
    /// able to send this number of bytes. In this case we will return
    /// `WriteZero` error however.
    pub fn size(&self) -> u64 {
        return self.size;
    }
    /// Returns a future which resolves to original socket when file has been
    /// written into a file
    pub fn write_into<D: Destination>(self, dest: D) -> WriteFile<O, D> {
        if self.cached {
            WriteFile(self.pool.clone(), WriteState::Mem(self, dest))
        } else {
            WriteFile(self.pool.clone(), WriteState::WaitWrite(self, dest))
        }
    }
    /// Get inner file opener
    pub fn get_inner(&self) -> &O {
        return &self.file;
    }
    /// Get mutlable reference to inner file opener
    pub fn get_mut(&mut self) -> &mut O {
        return &mut self.file;
    }
}

impl<F: FileOpener, D: Destination> Future for WriteFile<F, D>
    where F: Send + 'static, D: Send + 'static,
{
    type Item = D;
    type Error = io::Error;
    fn poll(&mut self) -> Poll<D, io::Error> {
        use self::WriteState::*;
        loop {
            let (newstate, cont) = match mem::replace(&mut self.1, Empty) {
                Mem(mut file, mut dest) => {
                    let need_switch = match file.file.from_cache() {
                        Some(Ok(slice)) => {
                            if (slice.len() as u64) < file.size {
                                return Err(io::Error::new(
                                    io::ErrorKind::WriteZero,
                                    "cached file truncated during writing"));
                            }
                            let target_slice = &slice[file.offset as usize..];
                            // Not sure why we can reach it, but it's safe
                            if target_slice.len() == 0 {
                                return Ok(Async::Ready(dest));
                            }
                            match dest.write(target_slice) {
                                Ok(0) => {
                                    return Err(io::Error::new(
                                        io::ErrorKind::WriteZero,
                                        "connection closed while sending \
                                         file from cache"));
                                }
                                Ok(bytes) => {
                                    file.offset += bytes as u64;
                                    if file.offset >= file.size {
                                        return Ok(Async::Ready(dest));
                                    }
                                }
                                Err(e) => {
                                    return Err(e);
                                }
                            }
                            false
                        }
                        Some(Err(e)) => {
                            return Err(e);
                        }
                        None => {
                            // File evicted from cache in the middle of sending
                            // Switch to non-cached variant
                            // TODO(tailhook) should we log it?
                            true
                        }
                    };
                    if need_switch {
                        (WaitWrite(file, dest), true)
                    } else {
                        (Mem(file, dest), false)
                    }
                }
                WaitSend(mut future) => {
                    match future.poll() {
                        Ok(Async::Ready((file, dest))) => {
                            if file.size <= file.offset {
                                return Ok(Async::Ready(dest));
                            } else {
                                (WaitWrite(file, dest), true)
                            }
                        }
                        Ok(Async::NotReady) => (WaitSend(future), false),
                        Err(e) => return Err(e),
                    }
                }
                WaitWrite(mut file, mut dest) => {
                    match dest.poll_write() {
                        Async::Ready(()) => {
                            (WaitSend(self.0.pool.spawn_fn(move || {
                                loop {
                                    match dest.write_file(&mut file) {
                                        Ok(0) => {
                                            return Err(io::Error::new(
                                                io::ErrorKind::WriteZero,
                                                "connection closed while \
                                                 sending a file"))
                                        }
                                        Ok(bytes_sent) => {
                                            file.offset += bytes_sent as u64;
                                            if file.offset >= file.size {
                                                return Ok((file, dest));
                                            } else {
                                                continue;
                                            }
                                        }
                                        Err(ref e)
                                        if e.kind() == io::ErrorKind::WouldBlock
                                        => {
                                            return Ok((file, dest))
                                        }
                                        Err(e) => return Err(e),
                                    }
                                }
                            })), true)
                        }
                        Async::NotReady => {
                            (WaitWrite(file, dest), false)
                        }
                    }
                }
                Empty => unreachable!(),
            };
            self.1 = newstate;
            if !cont {
                return Ok(Async::NotReady);
            }
        }
    }
}

#[cfg(unix)]
impl FileReader for File {
    fn read_at(&self, offset: u64, buf: &mut [u8]) -> io::Result<usize> {
        use libc::{pread, c_void};
        use std::os::unix::io::AsRawFd;

        let rc = unsafe { pread(self.as_raw_fd(),
            buf.as_ptr() as *mut c_void,
            buf.len(), offset as i64) };
        if rc < 0 {
            Err(io::Error::last_os_error())
        } else {
            Ok(rc as usize)
        }
    }
}

#[cfg(windows)]
impl FileReader for Mutex<File> {
    fn read_at(&self, offset: u64, buf: &mut [u8]) -> io::Result<usize> {
        use std::io::{Read, Seek};

        let mut real_file = self.lock().expect("mutex is not poisoned");
        real_file.seek(io::SeekFrom::Start(offset))?;
        real_file.read(buf)
    }
}