src/server/worker/api_server.go
func (a *APIServer) processDatums(pachClient *client.APIClient, logger *taggedLogger, jobInfo *pps.JobInfo,
	df DatumFactory, low, high int64, skip map[string]struct{}, useParentHashTree bool) (result *processResult, retErr error) {
	defer func() {
		if err := a.datumCache.Clear(); err != nil && retErr == nil {
			logger.Logf("error clearing datum cache: %v", err)
		}
		if err := a.datumStatsCache.Clear(); err != nil && retErr == nil {
			logger.Logf("error clearing datum stats cache: %v", err)
		}
	}()
	ctx := pachClient.Ctx()
	objClient, err := obj.NewClientFromEnv(a.hashtreeStorage)
	if err != nil {
		return nil, err
	}
	stats := &pps.ProcessStats{}
	var statsMu sync.Mutex
	result = &processResult{}
	var eg errgroup.Group
	limiter := limit.New(int(a.pipelineInfo.MaxQueueSize))
	for i := low; i < high; i++ {
		datumIdx := i

		limiter.Acquire()
		atomic.AddInt64(&a.queueSize, 1)
		eg.Go(func() (retErr error) {
			defer limiter.Release()
			defer atomic.AddInt64(&a.queueSize, -1)

			data := df.Datum(int(datumIdx))
			logger, err := a.getTaggedLogger(pachClient, jobInfo.Job.ID, data, a.pipelineInfo.EnableStats)
			if err != nil {
				return err
			}
			// Hash inputs
			tag := HashDatum(a.pipelineInfo.Pipeline.Name, a.pipelineInfo.Salt, data)
			if _, ok := skip[tag]; ok {
				if !useParentHashTree {
					if err := a.cacheHashtree(pachClient, tag, datumIdx); err != nil {
						return err
					}
				}
				atomic.AddInt64(&result.datumsSkipped, 1)
				logger.Logf("skipping datum")
				return nil
			}
			if _, err := pachClient.InspectTag(ctx, client.NewTag(tag)); err == nil {
				if err := a.cacheHashtree(pachClient, tag, datumIdx); err != nil {
					return err
				}
				atomic.AddInt64(&result.datumsSkipped, 1)
				logger.Logf("skipping datum")
				return nil
			}
			subStats := &pps.ProcessStats{}
			var inputTree, outputTree *hashtree.Ordered
			var statsTree *hashtree.Unordered
			if a.pipelineInfo.EnableStats {
				statsRoot := path.Join("/", logger.template.DatumID)
				inputTree = hashtree.NewOrdered(path.Join(statsRoot, "pfs"))
				outputTree = hashtree.NewOrdered(path.Join(statsRoot, "pfs", "out"))
				statsTree = hashtree.NewUnordered(statsRoot)
				// Write job id to stats tree
				statsTree.PutFile(fmt.Sprintf("job:%s", jobInfo.Job.ID), nil, 0)
				// Write index in datum factory to stats tree
				object, size, err := pachClient.PutObject(strings.NewReader(fmt.Sprint(int(datumIdx))))
				if err != nil {
					return err
				}
				objectInfo, err := pachClient.InspectObject(object.Hash)
				if err != nil {
					return err
				}
				h, err := pfs.DecodeHash(object.Hash)
				if err != nil {
					return err
				}
				statsTree.PutFile("index", h, size, objectInfo.BlockRef)
				defer func() {
					if err := a.writeStats(pachClient, objClient, tag, subStats, logger, inputTree, outputTree, statsTree, datumIdx); err != nil && retErr == nil {
						retErr = err
					}
				}()
			}

			env := a.userCodeEnv(jobInfo.Job.ID, jobInfo.OutputCommit.ID, data)
			var dir string
			var failures int64
			if err := backoff.RetryNotify(func() error {
				if isDone(ctx) {
					return ctx.Err() // timeout or cancelled job--don't run datum
				}
				// Download input data
				puller := filesync.NewPuller()
				// TODO parent tag shouldn't be nil
				var err error
				dir, err = a.downloadData(pachClient, logger, data, puller, subStats, inputTree)
				// We run these cleanup functions no matter what, so that if
				// downloadData partially succeeded, we still clean up the resources.
				defer func() {
					if err := os.RemoveAll(dir); err != nil && retErr == nil {
						retErr = err
					}
				}()
				// It's important that we run puller.CleanUp before os.RemoveAll,
				// because otherwise puller.Cleanup might try tp open pipes that have
				// been deleted.
				defer func() {
					if _, err := puller.CleanUp(); err != nil && retErr == nil {
						retErr = err
					}
				}()
				if err != nil {
					return fmt.Errorf("error downloadData: %v", err)
				}
				a.runMu.Lock()
				defer a.runMu.Unlock()
				// shadow ctx and pachClient for the context of processing this one datum
				ctx, cancel := context.WithCancel(ctx)
				pachClient := pachClient.WithCtx(ctx)
				func() {
					a.statusMu.Lock()
					defer a.statusMu.Unlock()
					a.jobID = jobInfo.Job.ID
					a.data = data
					a.started = time.Now()
					a.cancel = cancel
					a.stats = stats
				}()
				if err := os.MkdirAll(client.PPSInputPrefix, 0777); err != nil {
					return err
				}
				if err := a.linkData(data, dir); err != nil {
					return fmt.Errorf("error linkData: %v", err)
				}
				defer func() {
					if err := a.unlinkData(data); err != nil && retErr == nil {
						retErr = fmt.Errorf("error unlinkData: %v", err)
					}
				}()
				// If the pipeline spec set a custom user to execute the
				// process, make sure `/pfs` and its content are owned by it
				if a.uid != nil && a.gid != nil {
					filepath.Walk("/pfs", func(name string, info os.FileInfo, err error) error {
						if err == nil {
							err = os.Chown(name, int(*a.uid), int(*a.gid))
						}
						return err
					})
				}
				if err := a.runUserCode(ctx, logger, env, subStats, jobInfo.DatumTimeout); err != nil {
					if a.pipelineInfo.Transform.ErrCmd != nil && failures == jobInfo.DatumTries-1 {
						if err = a.runUserErrorHandlingCode(ctx, logger, env, subStats, jobInfo.DatumTimeout); err != nil {
							return fmt.Errorf("error runUserErrorHandlingCode: %v", err)
						}
						return errDatumRecovered
					}
					return fmt.Errorf("error runUserCode: %v", err)
				}
				// CleanUp is idempotent so we can call it however many times we want.
				// The reason we are calling it here is that the puller could've
				// encountered an error as it was lazily loading files, in which case
				// the output might be invalid since as far as the user's code is
				// concerned, they might've just seen an empty or partially completed
				// file.
				downSize, err := puller.CleanUp()
				if err != nil {
					logger.Logf("puller encountered an error while cleaning up: %+v", err)
					return err
				}
				atomic.AddUint64(&subStats.DownloadBytes, uint64(downSize))
				a.reportDownloadSizeStats(float64(downSize), logger)
				return a.uploadOutput(pachClient, dir, tag, logger, data, subStats, outputTree, datumIdx)
			}, &backoff.ZeroBackOff{}, func(err error, d time.Duration) error {
				if isDone(ctx) {
					return ctx.Err() // timeout or cancelled job, err out and don't retry
				}
				failures++
				if failures >= jobInfo.DatumTries {
					logger.Logf("failed to process datum with error: %+v", err)
					if statsTree != nil {
						object, size, err := pachClient.PutObject(strings.NewReader(err.Error()))
						if err != nil {
							logger.stderrLog.Printf("could not put error object: %s\n", err)
						} else {
							objectInfo, err := pachClient.InspectObject(object.Hash)
							if err != nil {
								return err
							}
							h, err := pfs.DecodeHash(object.Hash)
							if err != nil {
								return err
							}
							statsTree.PutFile("failure", h, size, objectInfo.BlockRef)
						}
					}
					return err
				}
				logger.Logf("failed processing datum: %v, retrying in %v", err, d)
				return nil
			}); err == errDatumRecovered {
				atomic.AddInt64(&result.datumsRecovered, 1)
				return nil
			} else if err != nil {
				result.failedDatumID = a.DatumID(data)
				atomic.AddInt64(&result.datumsFailed, 1)
				return nil
			}
			statsMu.Lock()
			defer statsMu.Unlock()
			if err := mergeStats(stats, subStats); err != nil {
				logger.Logf("failed to merge Stats: %v", err)
			}
			return nil
		})
	}
	if err := eg.Wait(); err != nil {
		return nil, err
	}
	if _, err := col.NewSTM(ctx, a.etcdClient, func(stm col.STM) error {
		jobs := a.jobs.ReadWrite(stm)
		jobID := jobInfo.Job.ID
		jobPtr := &pps.EtcdJobInfo{}
		if err := jobs.Get(jobID, jobPtr); err != nil {
			return err
		}
		if jobPtr.Stats == nil {
			jobPtr.Stats = &pps.ProcessStats{}
		}
		if err := mergeStats(jobPtr.Stats, stats); err != nil {
			logger.Logf("failed to merge Stats: %v", err)
		}
		return jobs.Put(jobID, jobPtr)
	}); err != nil {
		return nil, err
	}
	result.datumsProcessed = high - low - result.datumsSkipped - result.datumsFailed - result.datumsRecovered
	// Merge datum hashtrees into a chunk hashtree, then cache it.
	if err := a.mergeChunk(logger, high, result); err != nil {
		return nil, err
	}
	return result, nil
}
func (a *APIServer) processDatums(pachClient *client.APIClient, logger *taggedLogger, jobInfo *pps.JobInfo,
	df DatumFactory, low, high int64, skip map[string]struct{}, useParentHashTree bool) (result *processResult, retErr error) {
	defer func() {
		if err := a.datumCache.Clear(); err != nil && retErr == nil {
			logger.Logf("error clearing datum cache: %v", err)
		}
		if err := a.datumStatsCache.Clear(); err != nil && retErr == nil {
			logger.Logf("error clearing datum stats cache: %v", err)
		}
	}()
	ctx := pachClient.Ctx()
	objClient, err := obj.NewClientFromEnv(a.hashtreeStorage)
	if err != nil {
		return nil, err
	}
	stats := &pps.ProcessStats{}
	var statsMu sync.Mutex
	result = &processResult{}
	var eg errgroup.Group
	limiter := limit.New(int(a.pipelineInfo.MaxQueueSize))
	for i := low; i < high; i++ {
		datumIdx := i

		limiter.Acquire()
		atomic.AddInt64(&a.queueSize, 1)
		eg.Go(func() (retErr error) {
			defer limiter.Release()
			defer atomic.AddInt64(&a.queueSize, -1)

			data := df.Datum(int(datumIdx))
			logger, err := a.getTaggedLogger(pachClient, jobInfo.Job.ID, data, a.pipelineInfo.EnableStats)
			if err != nil {
				return err
			}
			// Hash inputs
			tag := HashDatum(a.pipelineInfo.Pipeline.Name, a.pipelineInfo.Salt, data)
			if _, ok := skip[tag]; ok {
				if !useParentHashTree {
					if err := a.cacheHashtree(pachClient, tag, datumIdx); err != nil {
						return err
					}
				}
				atomic.AddInt64(&result.datumsSkipped, 1)
				logger.Logf("skipping datum")
				return nil
			}
			if _, err := pachClient.InspectTag(ctx, client.NewTag(tag)); err == nil {
				if err := a.cacheHashtree(pachClient, tag, datumIdx); err != nil {
					return err
				}
				atomic.AddInt64(&result.datumsSkipped, 1)
				logger.Logf("skipping datum")
				return nil
			}
			subStats := &pps.ProcessStats{}
			var inputTree, outputTree *hashtree.Ordered
			var statsTree *hashtree.Unordered
			if a.pipelineInfo.EnableStats {
				statsRoot := path.Join("/", logger.template.DatumID)
				inputTree = hashtree.NewOrdered(path.Join(statsRoot, "pfs"))
				outputTree = hashtree.NewOrdered(path.Join(statsRoot, "pfs", "out"))
				statsTree = hashtree.NewUnordered(statsRoot)
				// Write job id to stats tree
				statsTree.PutFile(fmt.Sprintf("job:%s", jobInfo.Job.ID), nil, 0)
				// Write index in datum factory to stats tree
				object, size, err := pachClient.PutObject(strings.NewReader(fmt.Sprint(int(datumIdx))))
				if err != nil {
					return err
				}
				objectInfo, err := pachClient.InspectObject(object.Hash)
				if err != nil {
					return err
				}
				h, err := pfs.DecodeHash(object.Hash)
				if err != nil {
					return err
				}
				statsTree.PutFile("index", h, size, objectInfo.BlockRef)
				defer func() {
					if err := a.writeStats(pachClient, objClient, tag, subStats, logger, inputTree, outputTree, statsTree, datumIdx); err != nil && retErr == nil {
						retErr = err
					}
				}()
			}

			env := a.userCodeEnv(jobInfo.Job.ID, jobInfo.OutputCommit.ID, data)
			var dir string
			var failures int64
			if err := backoff.RetryNotify(func() error {
				if isDone(ctx) {
					return ctx.Err() // timeout or cancelled job--don't run datum
				}
				// Download input data
				puller := filesync.NewPuller()
				// TODO parent tag shouldn't be nil
				var err error
				dir, err = a.downloadData(pachClient, logger, data, puller, subStats, inputTree)
				// We run these cleanup functions no matter what, so that if
				// downloadData partially succeeded, we still clean up the resources.
				defer func() {
					if err := os.RemoveAll(dir); err != nil && retErr == nil {
						retErr = err
					}
				}()
				// It's important that we run puller.CleanUp before os.RemoveAll,
				// because otherwise puller.Cleanup might try tp open pipes that have
				// been deleted.
				defer func() {
					if _, err := puller.CleanUp(); err != nil && retErr == nil {
						retErr = err
					}
				}()
				if err != nil {
					return fmt.Errorf("error downloadData: %v", err)
				}
				a.runMu.Lock()
				defer a.runMu.Unlock()
				// shadow ctx and pachClient for the context of processing this one datum
				ctx, cancel := context.WithCancel(ctx)
				pachClient := pachClient.WithCtx(ctx)
				func() {
					a.statusMu.Lock()
					defer a.statusMu.Unlock()
					a.jobID = jobInfo.Job.ID
					a.data = data
					a.started = time.Now()
					a.cancel = cancel
					a.stats = stats
				}()
				if err := os.MkdirAll(client.PPSInputPrefix, 0777); err != nil {
					return err
				}
				if err := a.linkData(data, dir); err != nil {
					return fmt.Errorf("error linkData: %v", err)
				}
				defer func() {
					if err := a.unlinkData(data); err != nil && retErr == nil {
						retErr = fmt.Errorf("error unlinkData: %v", err)
					}
				}()
				// If the pipeline spec set a custom user to execute the
				// process, make sure `/pfs` and its content are owned by it
				if a.uid != nil && a.gid != nil {
					filepath.Walk("/pfs", func(name string, info os.FileInfo, err error) error {
						if err == nil {
							err = os.Chown(name, int(*a.uid), int(*a.gid))
						}
						return err
					})
				}
				if err := a.runUserCode(ctx, logger, env, subStats, jobInfo.DatumTimeout); err != nil {
					if a.pipelineInfo.Transform.ErrCmd != nil && failures == jobInfo.DatumTries-1 {
						if err = a.runUserErrorHandlingCode(ctx, logger, env, subStats, jobInfo.DatumTimeout); err != nil {
							return fmt.Errorf("error runUserErrorHandlingCode: %v", err)
						}
						return errDatumRecovered
					}
					return fmt.Errorf("error runUserCode: %v", err)
				}
				// CleanUp is idempotent so we can call it however many times we want.
				// The reason we are calling it here is that the puller could've
				// encountered an error as it was lazily loading files, in which case
				// the output might be invalid since as far as the user's code is
				// concerned, they might've just seen an empty or partially completed
				// file.
				downSize, err := puller.CleanUp()
				if err != nil {
					logger.Logf("puller encountered an error while cleaning up: %+v", err)
					return err
				}
				atomic.AddUint64(&subStats.DownloadBytes, uint64(downSize))
				a.reportDownloadSizeStats(float64(downSize), logger)
				return a.uploadOutput(pachClient, dir, tag, logger, data, subStats, outputTree, datumIdx)
			}, &backoff.ZeroBackOff{}, func(err error, d time.Duration) error {
				if isDone(ctx) {
					return ctx.Err() // timeout or cancelled job, err out and don't retry
				}
				failures++
				if failures >= jobInfo.DatumTries {
					logger.Logf("failed to process datum with error: %+v", err)
					if statsTree != nil {
						object, size, err := pachClient.PutObject(strings.NewReader(err.Error()))
						if err != nil {
							logger.stderrLog.Printf("could not put error object: %s\n", err)
						} else {
							objectInfo, err := pachClient.InspectObject(object.Hash)
							if err != nil {
								return err
							}
							h, err := pfs.DecodeHash(object.Hash)
							if err != nil {
								return err
							}
							statsTree.PutFile("failure", h, size, objectInfo.BlockRef)
						}
					}
					return err
				}
				logger.Logf("failed processing datum: %v, retrying in %v", err, d)
				return nil
			}); err == errDatumRecovered {
				atomic.AddInt64(&result.datumsRecovered, 1)
				return nil
			} else if err != nil {
				result.failedDatumID = a.DatumID(data)
				atomic.AddInt64(&result.datumsFailed, 1)
				return nil
			}
			statsMu.Lock()
			defer statsMu.Unlock()
			if err := mergeStats(stats, subStats); err != nil {
				logger.Logf("failed to merge Stats: %v", err)
			}
			return nil
		})
	}
	if err := eg.Wait(); err != nil {
		return nil, err
	}
	if _, err := col.NewSTM(ctx, a.etcdClient, func(stm col.STM) error {
		jobs := a.jobs.ReadWrite(stm)
		jobID := jobInfo.Job.ID
		jobPtr := &pps.EtcdJobInfo{}
		if err := jobs.Get(jobID, jobPtr); err != nil {
			return err
		}
		if jobPtr.Stats == nil {
			jobPtr.Stats = &pps.ProcessStats{}
		}
		if err := mergeStats(jobPtr.Stats, stats); err != nil {
			logger.Logf("failed to merge Stats: %v", err)
		}
		return jobs.Put(jobID, jobPtr)
	}); err != nil {
		return nil, err
	}
	result.datumsProcessed = high - low - result.datumsSkipped - result.datumsFailed - result.datumsRecovered
	// Merge datum hashtrees into a chunk hashtree, then cache it.
	if err := a.mergeChunk(logger, high, result); err != nil {
		return nil, err
	}
	return result, nil
}
// processDatums processes datums from low to high in df, if a datum fails it
// returns the id of the failed datum it also may return a variety of errors
// such as network errors.