github.com/pachyderm/pachyderm@v1.13.4/src/server/worker/pipeline/transform/status.go (about)

     1  package transform
     2  
     3  import (
     4  	"sync"
     5  	"sync/atomic"
     6  	"time"
     7  
     8  	"github.com/gogo/protobuf/types"
     9  
    10  	"github.com/pachyderm/pachyderm/src/client/pps"
    11  	"github.com/pachyderm/pachyderm/src/server/worker/common"
    12  )
    13  
    14  // Status is a struct representing the current status of the transform worker,
    15  // its public interface only allows getting the status of a task and canceling
    16  // the currently-processing datum.
    17  type Status struct {
    18  	mutex         sync.Mutex
    19  	jobID         string
    20  	stats         *pps.ProcessStats
    21  	queueSize     *int64
    22  	dataProcessed *int64
    23  	dataRecovered *int64
    24  	datum         []*pps.InputFile
    25  	cancel        func()
    26  	started       time.Time
    27  }
    28  
    29  func convertInputs(inputs []*common.Input) []*pps.InputFile {
    30  	var result []*pps.InputFile
    31  	for _, input := range inputs {
    32  		result = append(result, &pps.InputFile{
    33  			Path: input.FileInfo.File.Path,
    34  			Hash: input.FileInfo.Hash,
    35  		})
    36  	}
    37  	return result
    38  }
    39  
    40  func (s *Status) withLock(cb func()) {
    41  	s.mutex.Lock()
    42  	cb()
    43  	s.mutex.Unlock()
    44  }
    45  
    46  func (s *Status) withJob(jobID string, cb func() error) error {
    47  	s.withLock(func() {
    48  		s.jobID = jobID
    49  	})
    50  
    51  	defer s.withLock(func() {
    52  		s.jobID = ""
    53  	})
    54  
    55  	return cb()
    56  }
    57  
    58  func (s *Status) withStats(stats *pps.ProcessStats, queueSize *int64, dataProcessed, dataRecovered *int64, cb func() error) error {
    59  	s.withLock(func() {
    60  		s.stats = stats
    61  		s.queueSize = queueSize
    62  		s.dataProcessed = dataProcessed
    63  		s.dataRecovered = dataRecovered
    64  	})
    65  
    66  	defer s.withLock(func() {
    67  		s.stats = nil
    68  		s.queueSize = nil
    69  		s.dataProcessed = nil
    70  		s.dataRecovered = nil
    71  	})
    72  
    73  	return cb()
    74  }
    75  
    76  func (s *Status) withDatum(inputs []*common.Input, cancel func(), cb func() error) error {
    77  	s.withLock(func() {
    78  		s.datum = convertInputs(inputs)
    79  		s.cancel = cancel
    80  		s.started = time.Now()
    81  	})
    82  
    83  	defer s.withLock(func() {
    84  		s.datum = nil
    85  		s.cancel = nil
    86  		s.started = time.Time{}
    87  	})
    88  
    89  	return cb()
    90  }
    91  
    92  // GetStatus returns the current WorkerStatus for the transform worker
    93  func (s *Status) GetStatus() (*pps.WorkerStatus, error) {
    94  	s.mutex.Lock()
    95  	defer s.mutex.Unlock()
    96  
    97  	started, err := types.TimestampProto(s.started)
    98  	if err != nil {
    99  		return nil, err
   100  	}
   101  	result := &pps.WorkerStatus{
   102  		JobID:   s.jobID,
   103  		Data:    s.datum,
   104  		Started: started,
   105  	}
   106  	if s.queueSize != nil {
   107  		result.QueueSize = atomic.LoadInt64(s.queueSize)
   108  	}
   109  	if s.dataProcessed != nil {
   110  		result.DataProcessed = atomic.LoadInt64(s.dataProcessed)
   111  	}
   112  	if s.dataRecovered != nil {
   113  		result.DataRecovered = atomic.LoadInt64(s.dataRecovered)
   114  	}
   115  	return result, nil
   116  }
   117  
   118  // Cancel cancels the currently running datum if it matches the specified job and inputs
   119  func (s *Status) Cancel(jobID string, datumFilter []string) bool {
   120  	s.mutex.Lock()
   121  	defer s.mutex.Unlock()
   122  
   123  	if jobID == s.jobID && common.MatchDatum(datumFilter, s.datum) {
   124  		// Fields will be cleared as the worker stack unwinds
   125  		s.cancel()
   126  		return true
   127  	}
   128  	return false
   129  }