github.com/kunnos/engine@v1.13.1/distribution/xfer/download_test.go (about) 1 package xfer 2 3 import ( 4 "bytes" 5 "errors" 6 "fmt" 7 "io" 8 "io/ioutil" 9 "runtime" 10 "sync/atomic" 11 "testing" 12 "time" 13 14 "github.com/docker/distribution" 15 "github.com/docker/distribution/digest" 16 "github.com/docker/docker/image" 17 "github.com/docker/docker/layer" 18 "github.com/docker/docker/pkg/progress" 19 "golang.org/x/net/context" 20 ) 21 22 const maxDownloadConcurrency = 3 23 24 type mockLayer struct { 25 layerData bytes.Buffer 26 diffID layer.DiffID 27 chainID layer.ChainID 28 parent layer.Layer 29 } 30 31 func (ml *mockLayer) TarStream() (io.ReadCloser, error) { 32 return ioutil.NopCloser(bytes.NewBuffer(ml.layerData.Bytes())), nil 33 } 34 35 func (ml *mockLayer) TarStreamFrom(layer.ChainID) (io.ReadCloser, error) { 36 return nil, fmt.Errorf("not implemented") 37 } 38 39 func (ml *mockLayer) ChainID() layer.ChainID { 40 return ml.chainID 41 } 42 43 func (ml *mockLayer) DiffID() layer.DiffID { 44 return ml.diffID 45 } 46 47 func (ml *mockLayer) Parent() layer.Layer { 48 return ml.parent 49 } 50 51 func (ml *mockLayer) Size() (size int64, err error) { 52 return 0, nil 53 } 54 55 func (ml *mockLayer) DiffSize() (size int64, err error) { 56 return 0, nil 57 } 58 59 func (ml *mockLayer) Metadata() (map[string]string, error) { 60 return make(map[string]string), nil 61 } 62 63 type mockLayerStore struct { 64 layers map[layer.ChainID]*mockLayer 65 } 66 67 func createChainIDFromParent(parent layer.ChainID, dgsts ...layer.DiffID) layer.ChainID { 68 if len(dgsts) == 0 { 69 return parent 70 } 71 if parent == "" { 72 return createChainIDFromParent(layer.ChainID(dgsts[0]), dgsts[1:]...) 73 } 74 // H = "H(n-1) SHA256(n)" 75 dgst := digest.FromBytes([]byte(string(parent) + " " + string(dgsts[0]))) 76 return createChainIDFromParent(layer.ChainID(dgst), dgsts[1:]...) 77 } 78 79 func (ls *mockLayerStore) Map() map[layer.ChainID]layer.Layer { 80 layers := map[layer.ChainID]layer.Layer{} 81 82 for k, v := range ls.layers { 83 layers[k] = v 84 } 85 86 return layers 87 } 88 89 func (ls *mockLayerStore) Register(reader io.Reader, parentID layer.ChainID) (layer.Layer, error) { 90 return ls.RegisterWithDescriptor(reader, parentID, distribution.Descriptor{}) 91 } 92 93 func (ls *mockLayerStore) RegisterWithDescriptor(reader io.Reader, parentID layer.ChainID, _ distribution.Descriptor) (layer.Layer, error) { 94 var ( 95 parent layer.Layer 96 err error 97 ) 98 99 if parentID != "" { 100 parent, err = ls.Get(parentID) 101 if err != nil { 102 return nil, err 103 } 104 } 105 106 l := &mockLayer{parent: parent} 107 _, err = l.layerData.ReadFrom(reader) 108 if err != nil { 109 return nil, err 110 } 111 l.diffID = layer.DiffID(digest.FromBytes(l.layerData.Bytes())) 112 l.chainID = createChainIDFromParent(parentID, l.diffID) 113 114 ls.layers[l.chainID] = l 115 return l, nil 116 } 117 118 func (ls *mockLayerStore) Get(chainID layer.ChainID) (layer.Layer, error) { 119 l, ok := ls.layers[chainID] 120 if !ok { 121 return nil, layer.ErrLayerDoesNotExist 122 } 123 return l, nil 124 } 125 126 func (ls *mockLayerStore) Release(l layer.Layer) ([]layer.Metadata, error) { 127 return []layer.Metadata{}, nil 128 } 129 func (ls *mockLayerStore) CreateRWLayer(string, layer.ChainID, string, layer.MountInit, map[string]string) (layer.RWLayer, error) { 130 return nil, errors.New("not implemented") 131 } 132 133 func (ls *mockLayerStore) GetRWLayer(string) (layer.RWLayer, error) { 134 return nil, errors.New("not implemented") 135 } 136 137 func (ls *mockLayerStore) ReleaseRWLayer(layer.RWLayer) ([]layer.Metadata, error) { 138 return nil, errors.New("not implemented") 139 } 140 func (ls *mockLayerStore) GetMountID(string) (string, error) { 141 return "", errors.New("not implemented") 142 } 143 144 func (ls *mockLayerStore) Cleanup() error { 145 return nil 146 } 147 148 func (ls *mockLayerStore) DriverStatus() [][2]string { 149 return [][2]string{} 150 } 151 152 func (ls *mockLayerStore) DriverName() string { 153 return "mock" 154 } 155 156 type mockDownloadDescriptor struct { 157 currentDownloads *int32 158 id string 159 diffID layer.DiffID 160 registeredDiffID layer.DiffID 161 expectedDiffID layer.DiffID 162 simulateRetries int 163 } 164 165 // Key returns the key used to deduplicate downloads. 166 func (d *mockDownloadDescriptor) Key() string { 167 return d.id 168 } 169 170 // ID returns the ID for display purposes. 171 func (d *mockDownloadDescriptor) ID() string { 172 return d.id 173 } 174 175 // DiffID should return the DiffID for this layer, or an error 176 // if it is unknown (for example, if it has not been downloaded 177 // before). 178 func (d *mockDownloadDescriptor) DiffID() (layer.DiffID, error) { 179 if d.diffID != "" { 180 return d.diffID, nil 181 } 182 return "", errors.New("no diffID available") 183 } 184 185 func (d *mockDownloadDescriptor) Registered(diffID layer.DiffID) { 186 d.registeredDiffID = diffID 187 } 188 189 func (d *mockDownloadDescriptor) mockTarStream() io.ReadCloser { 190 // The mock implementation returns the ID repeated 5 times as a tar 191 // stream instead of actual tar data. The data is ignored except for 192 // computing IDs. 193 return ioutil.NopCloser(bytes.NewBuffer([]byte(d.id + d.id + d.id + d.id + d.id))) 194 } 195 196 // Download is called to perform the download. 197 func (d *mockDownloadDescriptor) Download(ctx context.Context, progressOutput progress.Output) (io.ReadCloser, int64, error) { 198 if d.currentDownloads != nil { 199 defer atomic.AddInt32(d.currentDownloads, -1) 200 201 if atomic.AddInt32(d.currentDownloads, 1) > maxDownloadConcurrency { 202 return nil, 0, errors.New("concurrency limit exceeded") 203 } 204 } 205 206 // Sleep a bit to simulate a time-consuming download. 207 for i := int64(0); i <= 10; i++ { 208 select { 209 case <-ctx.Done(): 210 return nil, 0, ctx.Err() 211 case <-time.After(10 * time.Millisecond): 212 progressOutput.WriteProgress(progress.Progress{ID: d.ID(), Action: "Downloading", Current: i, Total: 10}) 213 } 214 } 215 216 if d.simulateRetries != 0 { 217 d.simulateRetries-- 218 return nil, 0, errors.New("simulating retry") 219 } 220 221 return d.mockTarStream(), 0, nil 222 } 223 224 func (d *mockDownloadDescriptor) Close() { 225 } 226 227 func downloadDescriptors(currentDownloads *int32) []DownloadDescriptor { 228 return []DownloadDescriptor{ 229 &mockDownloadDescriptor{ 230 currentDownloads: currentDownloads, 231 id: "id1", 232 expectedDiffID: layer.DiffID("sha256:68e2c75dc5c78ea9240689c60d7599766c213ae210434c53af18470ae8c53ec1"), 233 }, 234 &mockDownloadDescriptor{ 235 currentDownloads: currentDownloads, 236 id: "id2", 237 expectedDiffID: layer.DiffID("sha256:64a636223116aa837973a5d9c2bdd17d9b204e4f95ac423e20e65dfbb3655473"), 238 }, 239 &mockDownloadDescriptor{ 240 currentDownloads: currentDownloads, 241 id: "id3", 242 expectedDiffID: layer.DiffID("sha256:58745a8bbd669c25213e9de578c4da5c8ee1c836b3581432c2b50e38a6753300"), 243 }, 244 &mockDownloadDescriptor{ 245 currentDownloads: currentDownloads, 246 id: "id2", 247 expectedDiffID: layer.DiffID("sha256:64a636223116aa837973a5d9c2bdd17d9b204e4f95ac423e20e65dfbb3655473"), 248 }, 249 &mockDownloadDescriptor{ 250 currentDownloads: currentDownloads, 251 id: "id4", 252 expectedDiffID: layer.DiffID("sha256:0dfb5b9577716cc173e95af7c10289322c29a6453a1718addc00c0c5b1330936"), 253 simulateRetries: 1, 254 }, 255 &mockDownloadDescriptor{ 256 currentDownloads: currentDownloads, 257 id: "id5", 258 expectedDiffID: layer.DiffID("sha256:0a5f25fa1acbc647f6112a6276735d0fa01e4ee2aa7ec33015e337350e1ea23d"), 259 }, 260 } 261 } 262 263 func TestSuccessfulDownload(t *testing.T) { 264 // TODO Windows: Fix this unit text 265 if runtime.GOOS == "windows" { 266 t.Skip("Needs fixing on Windows") 267 } 268 layerStore := &mockLayerStore{make(map[layer.ChainID]*mockLayer)} 269 ldm := NewLayerDownloadManager(layerStore, maxDownloadConcurrency) 270 271 progressChan := make(chan progress.Progress) 272 progressDone := make(chan struct{}) 273 receivedProgress := make(map[string]progress.Progress) 274 275 go func() { 276 for p := range progressChan { 277 receivedProgress[p.ID] = p 278 } 279 close(progressDone) 280 }() 281 282 var currentDownloads int32 283 descriptors := downloadDescriptors(¤tDownloads) 284 285 firstDescriptor := descriptors[0].(*mockDownloadDescriptor) 286 287 // Pre-register the first layer to simulate an already-existing layer 288 l, err := layerStore.Register(firstDescriptor.mockTarStream(), "") 289 if err != nil { 290 t.Fatal(err) 291 } 292 firstDescriptor.diffID = l.DiffID() 293 294 rootFS, releaseFunc, err := ldm.Download(context.Background(), *image.NewRootFS(), descriptors, progress.ChanOutput(progressChan)) 295 if err != nil { 296 t.Fatalf("download error: %v", err) 297 } 298 299 releaseFunc() 300 301 close(progressChan) 302 <-progressDone 303 304 if len(rootFS.DiffIDs) != len(descriptors) { 305 t.Fatal("got wrong number of diffIDs in rootfs") 306 } 307 308 for i, d := range descriptors { 309 descriptor := d.(*mockDownloadDescriptor) 310 311 if descriptor.diffID != "" { 312 if receivedProgress[d.ID()].Action != "Already exists" { 313 t.Fatalf("did not get 'Already exists' message for %v", d.ID()) 314 } 315 } else if receivedProgress[d.ID()].Action != "Pull complete" { 316 t.Fatalf("did not get 'Pull complete' message for %v", d.ID()) 317 } 318 319 if rootFS.DiffIDs[i] != descriptor.expectedDiffID { 320 t.Fatalf("rootFS item %d has the wrong diffID (expected: %v got: %v)", i, descriptor.expectedDiffID, rootFS.DiffIDs[i]) 321 } 322 323 if descriptor.diffID == "" && descriptor.registeredDiffID != rootFS.DiffIDs[i] { 324 t.Fatal("diffID mismatch between rootFS and Registered callback") 325 } 326 } 327 } 328 329 func TestCancelledDownload(t *testing.T) { 330 ldm := NewLayerDownloadManager(&mockLayerStore{make(map[layer.ChainID]*mockLayer)}, maxDownloadConcurrency) 331 332 progressChan := make(chan progress.Progress) 333 progressDone := make(chan struct{}) 334 335 go func() { 336 for range progressChan { 337 } 338 close(progressDone) 339 }() 340 341 ctx, cancel := context.WithCancel(context.Background()) 342 343 go func() { 344 <-time.After(time.Millisecond) 345 cancel() 346 }() 347 348 descriptors := downloadDescriptors(nil) 349 _, _, err := ldm.Download(ctx, *image.NewRootFS(), descriptors, progress.ChanOutput(progressChan)) 350 if err != context.Canceled { 351 t.Fatal("expected download to be cancelled") 352 } 353 354 close(progressChan) 355 <-progressDone 356 }