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