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         spans     map[metainfo.InfoHash]mmap_span.MMapSpan
  19         completed map[metainfo.InfoHash]bool
  20 }
  21
  22 func NewMMap(baseDir string) I {
  23         return &mmapStorage{
  24                 baseDir: baseDir,
  25         }
  26 }
  27
  28 func (me *mmapStorage) lazySpan(info *metainfo.InfoEx) error {
  29         if me.spans == nil {
  30                 me.spans = make(map[metainfo.InfoHash]mmap_span.MMapSpan)
  31         }
  32         if _, ok := me.spans[*info.Hash]; ok {
  33                 return nil
  34         }
  35         span, err := MMapTorrent(&info.Info, me.baseDir)
  36         if err != nil {
  37                 return err
  38         }
  39         me.spans[*info.Hash] = span
  40         return nil
  41 }
  42
  43 func (me *mmapStorage) Piece(p metainfo.Piece) Piece {
  44         err := me.lazySpan(p.Info)
  45         if err != nil {
  46                 panic(err)
  47         }
  48         return mmapStoragePiece{
  49                 storage:  me,
  50                 p:        p,
  51                 ReaderAt: io.NewSectionReader(me.spans[*p.Info.Hash], p.Offset(), p.Length()),
  52                 WriterAt: missinggo.NewSectionWriter(me.spans[*p.Info.Hash], p.Offset(), p.Length()),
  53         }
  54 }
  55
  56 type mmapStoragePiece struct {
  57         storage *mmapStorage
  58         p       metainfo.Piece
  59         io.ReaderAt
  60         io.WriterAt
  61 }
  62
  63 func (me mmapStoragePiece) GetIsComplete() bool {
  64         return me.storage.completed[me.p.Hash()]
  65 }
  66
  67 func (me mmapStoragePiece) MarkComplete() error {
  68         if me.storage.completed == nil {
  69                 me.storage.completed = make(map[metainfo.InfoHash]bool)
  70         }
  71         me.storage.completed[me.p.Hash()] = true
  72         return nil
  73 }
  74
  75 func MMapTorrent(md *metainfo.Info, location string) (mms mmap_span.MMapSpan, err error) {
  76         defer func() {
  77                 if err != nil {
  78                         mms.Close()
  79                 }
  80         }()
  81         for _, miFile := range md.UpvertedFiles() {
  82                 fileName := filepath.Join(append([]string{location, md.Name}, miFile.Path...)...)
  83                 err = os.MkdirAll(filepath.Dir(fileName), 0777)
  84                 if err != nil {
  85                         err = fmt.Errorf("error creating data directory %q: %s", filepath.Dir(fileName), err)
  86                         return
  87                 }
  88                 var file *os.File
  89                 file, err = os.OpenFile(fileName, os.O_CREATE|os.O_RDWR, 0666)
  90                 if err != nil {
  91                         return
  92                 }
  93                 func() {
  94                         defer file.Close()
  95                         var fi os.FileInfo
  96                         fi, err = file.Stat()
  97                         if err != nil {
  98                                 return
  99                         }
 100                         if fi.Size() < miFile.Length {
 101                                 err = file.Truncate(miFile.Length)
 102                                 if err != nil {
 103                                         return
 104                                 }
 105                         }
 106                         if miFile.Length == 0 {
 107                                 // Can't mmap() regions with length 0.
 108                                 return
 109                         }
 110                         var mMap mmap.MMap
 111                         mMap, err = mmap.MapRegion(file,
 112                                 int(miFile.Length), // Probably not great on <64 bit systems.
 113                                 mmap.RDWR, 0, 0)
 114                         if err != nil {
 115                                 err = fmt.Errorf("error mapping file %q, length %d: %s", file.Name(), miFile.Length, err)
 116                                 return
 117                         }
 118                         if int64(len(mMap)) != miFile.Length {
 119                                 panic("mmap has wrong length")
 120                         }
 121                         mms.Append(mMap)
 122                 }()
 123                 if err != nil {
 124                         return
 125                 }
 126         }
 127         return
 128 }