github.com/NVIDIA/aistore@v1.3.23-0.20240517131212-7df6609be51d/ext/dload/infostore.go (about)

     1  // Package dload implements functionality to download resources into AIS cluster from external source.
     2  /*
     3   * Copyright (c) 2018-2024, NVIDIA CORPORATION. All rights reserved.
     4   */
     5  package dload
     6  
     7  import (
     8  	"errors"
     9  	"net/http"
    10  	"sync"
    11  	"time"
    12  
    13  	"github.com/NVIDIA/aistore/cmn/cos"
    14  	"github.com/NVIDIA/aistore/cmn/debug"
    15  	"github.com/NVIDIA/aistore/cmn/kvdb"
    16  	"github.com/NVIDIA/aistore/core"
    17  	"github.com/NVIDIA/aistore/hk"
    18  )
    19  
    20  // TODO: stored only in memory, should be persisted at some point (powercycle)
    21  type infoStore struct {
    22  	*downloaderDB
    23  	dljobs map[string]*dljob
    24  	sync.RWMutex
    25  }
    26  
    27  func newInfoStore(driver kvdb.Driver) *infoStore {
    28  	db := newDownloadDB(driver)
    29  	is := &infoStore{
    30  		downloaderDB: db,
    31  		dljobs:       make(map[string]*dljob),
    32  	}
    33  	hk.Reg("downloader"+hk.NameSuffix, is.housekeep, hk.DayInterval)
    34  	return is
    35  }
    36  
    37  func (is *infoStore) getJob(id string) (*dljob, error) {
    38  	is.RLock()
    39  	defer is.RUnlock()
    40  
    41  	if ji, ok := is.dljobs[id]; ok {
    42  		return ji, nil
    43  	}
    44  	return nil, errJobNotFound
    45  }
    46  
    47  func (is *infoStore) getList(req *request) (jobs []*dljob) {
    48  	is.RLock()
    49  	for _, job := range is.dljobs {
    50  		if req.onlyActive && !_isRunning(job.finishedTime.Load()) {
    51  			continue
    52  		}
    53  		if req.regex == nil || req.regex.MatchString(job.description) {
    54  			jobs = append(jobs, job)
    55  		}
    56  	}
    57  	is.RUnlock()
    58  	return
    59  }
    60  
    61  func (is *infoStore) setJob(job jobif) (njob *dljob) {
    62  	njob = &dljob{
    63  		id:          job.ID(),
    64  		xid:         job.XactID(),
    65  		total:       job.Len(),
    66  		description: job.Description(),
    67  		startedTime: time.Now(),
    68  	}
    69  	is.Lock()
    70  	is.dljobs[job.ID()] = njob
    71  	is.Unlock()
    72  	return
    73  }
    74  
    75  func (is *infoStore) incFinished(id string) {
    76  	dljob, err := is.getJob(id)
    77  	debug.AssertNoErr(err)
    78  	dljob.finishedCnt.Inc()
    79  }
    80  
    81  func (is *infoStore) incSkipped(id string) {
    82  	dljob, err := is.getJob(id)
    83  	debug.AssertNoErr(err)
    84  	dljob.skippedCnt.Inc()
    85  	dljob.finishedCnt.Inc()
    86  }
    87  
    88  func (is *infoStore) incScheduled(id string) {
    89  	dljob, err := is.getJob(id)
    90  	debug.AssertNoErr(err)
    91  	dljob.scheduledCnt.Inc()
    92  }
    93  
    94  func (is *infoStore) incErrorCnt(id string) {
    95  	dljob, err := is.getJob(id)
    96  	debug.AssertNoErr(err)
    97  	dljob.errorCnt.Inc()
    98  }
    99  
   100  func (is *infoStore) setAllDispatched(id string, dispatched bool) {
   101  	dljob, err := is.getJob(id)
   102  	debug.AssertNoErr(err)
   103  	dljob.allDispatched.Store(dispatched)
   104  }
   105  
   106  func (is *infoStore) markFinished(id string) (error, bool /*aborted*/) {
   107  	dljob, err := is.getJob(id)
   108  	if err != nil {
   109  		debug.AssertNoErr(err)
   110  		return err, false
   111  	}
   112  	dljob.finishedTime.Store(time.Now())
   113  	return dljob.valid(), dljob.aborted.Load()
   114  }
   115  
   116  func (is *infoStore) setAborted(id string) {
   117  	dljob, err := is.getJob(id)
   118  	debug.AssertNoErr(err)
   119  	dljob.aborted.Store(true)
   120  	// NOTE: Don't set `FinishedTime` yet as we are not fully done.
   121  	//       The job now can be removed but there's no guarantee
   122  	//       that all tasks have been stopped and all resources were freed.
   123  }
   124  
   125  func (is *infoStore) delJob(id string) {
   126  	delete(is.dljobs, id)
   127  	is.downloaderDB.delete(id)
   128  }
   129  
   130  func (is *infoStore) housekeep() time.Duration {
   131  	const interval = hk.DayInterval
   132  
   133  	is.Lock()
   134  	for id, dljob := range is.dljobs {
   135  		if time.Since(dljob.finishedTime.Load()) > interval {
   136  			is.delJob(id)
   137  		}
   138  	}
   139  	is.Unlock()
   140  
   141  	return interval
   142  }
   143  
   144  func (is *infoStore) checkExists(req *request) (dljob *dljob, err error) {
   145  	dljob, err = is.getJob(req.id)
   146  	if err != nil {
   147  		debug.Assert(errors.Is(err, errJobNotFound))
   148  		err = cos.NewErrNotFound(core.T, "download job "+req.id)
   149  		req.errRsp(err, http.StatusNotFound)
   150  	}
   151  	return
   152  }