From 9549d98a7caeaf04a09e64bb16fd057ab8eb4fee Mon Sep 17 00:00:00 2001 From: miraclesu Date: Tue, 26 Mar 2013 17:22:07 +0800 Subject: [PATCH] Close friendly --- worker/agent.go | 6 ++++-- worker/worker.go | 8 +++++--- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/worker/agent.go b/worker/agent.go index 9c8cc01..1e17584 100644 --- a/worker/agent.go +++ b/worker/agent.go @@ -41,7 +41,7 @@ func newAgent(addr string, worker *Worker) (a *agent, err error) { func (a *agent) outLoop() { ok := true var job *Job - for ok { + for a.worker.running && ok { if job, ok = <-a.out; ok { if err := a.write(job.Encode()); err != nil { a.worker.err(err) @@ -92,7 +92,9 @@ func (a *agent) inLoop() { a.WriteJob(newJob(common.REQ, common.GRAB_JOB_UNIQ, nil)) case common.ERROR, common.ECHO_RES, common.JOB_ASSIGN_UNIQ, common.JOB_ASSIGN: job.agent = a - a.worker.in <- job + if a.worker.running { + a.worker.in <- job + } } } } diff --git a/worker/worker.go b/worker/worker.go index 33f6140..5f26b46 100644 --- a/worker/worker.go +++ b/worker/worker.go @@ -165,7 +165,7 @@ func (worker *Worker) removeFunc(funcname string) { func (worker *Worker) dealJob(job *Job) { defer func() { job.Close() - if worker.limit != nil { + if worker.running && worker.limit != nil { worker.limit <- true } }() @@ -185,7 +185,6 @@ func (worker *Worker) dealJob(job *Job) { // Main loop func (worker *Worker) Work() { defer func() { - worker.running = false for _, v := range worker.agents { v.Close() } @@ -220,6 +219,7 @@ func (worker *Worker) handleJob(job *Job) { // Close. func (worker *Worker) Close() { + worker.running = false close(worker.in) if worker.limit != nil { close(worker.limit) @@ -299,7 +299,9 @@ func (worker *Worker) exec(job *Job) (err error) { job.magicCode = common.REQ job.DataType = datatype job.Data = r.data - job.agent.WriteJob(job) + if worker.running { + job.agent.WriteJob(job) + } return }