storage/mmap.go

   1 package storage
   2
   3 import (
   4         "fmt"
   5         "io"
   6         "os"
   7         "path/filepath"
   8
   9         "github.com/anacrolix/missinggo"
  10         "github.com/edsrzf/mmap-go"
  11
  12         "github.com/anacrolix/torrent/metainfo"
  13         "github.com/anacrolix/torrent/mmap_span"
  14 )
  15
  16 type mmapStorage struct {
  17         baseDir    string
  18         completion pieceCompletion
  19 }
  20
  21 func NewMMap(baseDir string) Client {
  22         return &mmapStorage{
  23                 baseDir:    baseDir,
  24                 completion: pieceCompletionForDir(baseDir),
  25         }
  26 }
  27
  28 func (s *mmapStorage) OpenTorrent(info *metainfo.InfoEx) (t Torrent, err error) {
  29         span, err := mMapTorrent(&info.Info, s.baseDir)
  30         t = &mmapTorrentStorage{
  31                 span: span,
  32                 pc:   s.completion,
  33         }
  34         return
  35 }
  36
  37 type mmapTorrentStorage struct {
  38         span mmap_span.MMapSpan
  39         pc   pieceCompletion
  40 }
  41
  42 func (ts *mmapTorrentStorage) Piece(p metainfo.Piece) Piece {
  43         return mmapStoragePiece{
  44                 pc:       ts.pc,
  45                 p:        p,
  46                 ReaderAt: io.NewSectionReader(ts.span, p.Offset(), p.Length()),
  47                 WriterAt: missinggo.NewSectionWriter(ts.span, p.Offset(), p.Length()),
  48         }
  49 }
  50
  51 func (ts *mmapTorrentStorage) Close() error {
  52         ts.pc.Close()
  53         return ts.span.Close()
  54 }
  55
  56 type mmapStoragePiece struct {
  57         pc pieceCompletion
  58         p  metainfo.Piece
  59         io.ReaderAt
  60         io.WriterAt
  61 }
  62
  63 func (sp mmapStoragePiece) GetIsComplete() bool {
  64         return sp.pc.Get(sp.p)
  65 }
  66
  67 func (sp mmapStoragePiece) MarkComplete() error {
  68         sp.pc.Set(sp.p, true)
  69         return nil
  70 }
  71
  72 func mMapTorrent(md *metainfo.Info, location string) (mms mmap_span.MMapSpan, err error) {
  73         defer func() {
  74                 if err != nil {
  75                         mms.Close()
  76                 }
  77         }()
  78         for _, miFile := range md.UpvertedFiles() {
  79                 fileName := filepath.Join(append([]string{location, md.Name}, miFile.Path...)...)
  80                 err = os.MkdirAll(filepath.Dir(fileName), 0777)
  81                 if err != nil {
  82                         err = fmt.Errorf("error creating data directory %q: %s", filepath.Dir(fileName), err)
  83                         return
  84                 }
  85                 var file *os.File
  86                 file, err = os.OpenFile(fileName, os.O_CREATE|os.O_RDWR, 0666)
  87                 if err != nil {
  88                         return
  89                 }
  90                 func() {
  91                         defer file.Close()
  92                         var fi os.FileInfo
  93                         fi, err = file.Stat()
  94                         if err != nil {
  95                                 return
  96                         }
  97                         if fi.Size() < miFile.Length {
  98                                 err = file.Truncate(miFile.Length)
  99                                 if err != nil {
 100                                         return
 101                                 }
 102                         }
 103                         if miFile.Length == 0 {
 104                                 // Can't mmap() regions with length 0.
 105                                 return
 106                         }
 107                         var mMap mmap.MMap
 108                         mMap, err = mmap.MapRegion(file,
 109                                 int(miFile.Length), // Probably not great on <64 bit systems.
 110                                 mmap.RDWR, 0, 0)
 111                         if err != nil {
 112                                 err = fmt.Errorf("error mapping file %q, length %d: %s", file.Name(), miFile.Length, err)
 113                                 return
 114                         }
 115                         if int64(len(mMap)) != miFile.Length {
 116                                 panic("mmap has wrong length")
 117                         }
 118                         mms.Append(mMap)
 119                 }()
 120                 if err != nil {
 121                         return
 122                 }
 123         }
 124         return
 125 }