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