my_inetpeer.c

来自「一个基于linux的TCP/IP协议栈的实现」· C语言 代码 · 共 324 行

int myinet_peer_minttl = 120 * HZ;
int myinet_peer_maxttl = 10 * 60 * HZ;
static kmem_cache_t *mypeer_cachep;

static struct inet_peer mypeer_fake_node = {
	.avl_left	= &mypeer_fake_node,
	.avl_right	= &mypeer_fake_node,
	.avl_height	= 0
};
#define mypeer_avl_empty ( &mypeer_fake_node )
static struct inet_peer *mypeer_root = mypeer_avl_empty;
static DEFINE_RWLOCK(mypeer_pool_lock);
#define PEER_MAXDEPTH 40

DEFINE_SPINLOCK( myinet_peer_unused_lock );
static struct inet_peer *myinet_peer_unused_head;
struct inet_peer **myinet_peer_unused_tailp = &myinet_peer_unused_head;

#define node_height(x) x->avl_height
static volatile int mypeer_total;
int myinet_peer_threshold = 65536 + 128;

static void mypeer_check_expire(unsigned long dummy)
{

}

static struct timer_list mypeer_periodic_timer =
	TIMER_INITIALIZER(mypeer_check_expire, 0, 0);

int myinet_peer_gc_mintime = 10 * HZ,
    myinet_peer_gc_maxtime = 120 * HZ;

#define mylookup(daddr) 									\
({														\
	struct inet_peer *u, **v;							\
	stackptr = stack;									\
	*stackptr++ = &mypeer_root;							\
	for (u = mypeer_root; u != mypeer_avl_empty; ) {	\
		if (daddr == u->v4daddr)						\
			break;										\
		if (daddr < u->v4daddr)							\
			v = &u->avl_left;							\
		else											\
			v = &u->avl_right;							\
		*stackptr++ = v;								\
		u = *v;											\
	}													\
	u;													\
})

#define mylink_to_pool(n)					\
do {										\
	n->avl_height = 1;						\
	n->avl_left = mypeer_avl_empty;			\
	n->avl_right = mypeer_avl_empty;		\
	**--stackptr = n;						\
	mypeer_avl_rebalance(stack, stackptr);	\
} while(0)

#define mylookup_rightempty(start)						\
({														\
	struct inet_peer *u, **v;							\
	*stackptr++ = &start->avl_left;						\
	v = &start->avl_left;								\
	for (u = *v; u->avl_right != mypeer_avl_empty; ) {	\
		v = &u->avl_right;								\
		*stackptr++ = v;								\
		u = *v;											\
	}													\
	u;													\
})

static inline void	myinet_putpeer(struct inet_peer *p)
{
	spin_lock_bh( &myinet_peer_unused_lock );
	if (atomic_dec_and_test(&p->refcnt)) {
		p->unused_prevp = myinet_peer_unused_tailp;
		p->unused_next = NULL;
		*myinet_peer_unused_tailp = p;
		myinet_peer_unused_tailp = &p->unused_next;
		p->dtime = jiffies;
	}
	spin_unlock_bh( &myinet_peer_unused_lock );
}

static void mypeer_avl_rebalance(struct inet_peer **stack[],
				struct inet_peer ***stackend)
{
	struct inet_peer **nodep, *node, *l, *r;
	int lh, rh;

	while( stackend > stack ){
		nodep = *--stackend;
		node = *nodep;
		l = node->avl_left;
		r = node->avl_right;
		lh = node_height(l);
		rh = node_height(r);
		if (lh > rh + 1) { /* l: RH+2 */
			struct inet_peer *ll, *lr, *lrl, *lrr;
			int lrh;
			ll = l->avl_left;
			lr = l->avl_right;
			lrh = node_height(lr);
			if (lrh <= node_height(ll)) {	/* ll: RH+1 */
				node->avl_left = lr;	/* lr: RH or RH+1 */
				node->avl_right = r;	/* r: RH */
				node->avl_height = lrh + 1; /* RH+1 or RH+2 */
				l->avl_left = ll;	/* ll: RH+1 */
				l->avl_right = node;	/* node: RH+1 or RH+2 */
				l->avl_height = node->avl_height + 1;
				*nodep = l;
			} else { /* ll: RH, lr: RH+1 */
				lrl = lr->avl_left;	/* lrl: RH or RH-1 */
				lrr = lr->avl_right;	/* lrr: RH or RH-1 */
				node->avl_left = lrr;	/* lrr: RH or RH-1 */
				node->avl_right = r;	/* r: RH */
				node->avl_height = rh + 1; /* node: RH+1 */
				l->avl_left = ll;	/* ll: RH */
				l->avl_right = lrl;	/* lrl: RH or RH-1 */
				l->avl_height = rh + 1;	/* l: RH+1 */
				lr->avl_left = l;	/* l: RH+1 */
				lr->avl_right = node;	/* node: RH+1 */
				lr->avl_height = rh + 2;
				*nodep = lr;
			}
		} else if (rh > lh + 1) { /* r: LH+2 */
			struct inet_peer *rr, *rl, *rlr, *rll;
			int rlh;
			rr = r->avl_right;
			rl = r->avl_left;
			rlh = node_height(rl);
			if (rlh <= node_height(rr)) {	/* rr: LH+1 */
				node->avl_right = rl;	/* rl: LH or LH+1 */
				node->avl_left = l;	/* l: LH */
				node->avl_height = rlh + 1; /* LH+1 or LH+2 */
				r->avl_right = rr;	/* rr: LH+1 */
				r->avl_left = node;	/* node: LH+1 or LH+2 */
				r->avl_height = node->avl_height + 1;
				*nodep = r;
			} else { /* rr: RH, rl: RH+1 */
				rlr = rl->avl_right;	/* rlr: LH or LH-1 */
				rll = rl->avl_left;	/* rll: LH or LH-1 */
				node->avl_right = rll;	/* rll: LH or LH-1 */
				node->avl_left = l;	/* l: LH */
				node->avl_height = lh + 1; /* node: LH+1 */
				r->avl_right = rr;	/* rr: LH */
				r->avl_left = rlr;	/* rlr: LH or LH-1 */
				r->avl_height = lh + 1;	/* r: LH+1 */
				rl->avl_right = r;	/* r: LH+1 */
				rl->avl_left = node;	/* node: LH+1 */
				rl->avl_height = lh + 2;
				*nodep = rl;
			}
		} else {
			node->avl_height = (lh > rh ? lh : rh) + 1;
		}
	}
}

static void myunlink_from_unused(struct inet_peer *p)
{
	spin_lock_bh( &myinet_peer_unused_lock );
	if( p->unused_prevp != NULL ){
		*p->unused_prevp = p->unused_next;
		if (p->unused_next != NULL)
			p->unused_next->unused_prevp = p->unused_prevp;
		else
			myinet_peer_unused_tailp = p->unused_prevp;
		p->unused_prevp = NULL;
	}
	spin_unlock_bh( &myinet_peer_unused_lock );
}

static void myunlink_from_pool(struct inet_peer *p)
{
	int do_free;
	do_free = 0;

	write_lock_bh( &mypeer_pool_lock );
	if (atomic_read(&p->refcnt) == 1) {
		struct inet_peer **stack[PEER_MAXDEPTH];
		struct inet_peer ***stackptr, ***delp;
		if (mylookup(p->v4daddr) != p)
			BUG();
		delp = stackptr - 1;
		if (p->avl_left == mypeer_avl_empty) {
			*delp[0] = p->avl_right;
			--stackptr;
		} else {
			struct inet_peer *t;
			t = mylookup_rightempty(p);
			BUG_ON(*stackptr[-1] != t);
			**--stackptr = t->avl_left;
			*delp[0] = t;
			t->avl_left = p->avl_left;
			t->avl_right = p->avl_right;
			t->avl_height = p->avl_height;
			BUG_ON(delp[1] != &p->avl_left);
			delp[1] = &t->avl_left;
		}
		mypeer_avl_rebalance(stack, stackptr);
		mypeer_total--;
		do_free = 1;
	}
	write_unlock_bh( &mypeer_pool_lock );

	if (do_free)
		kmem_cache_free( mypeer_cachep, p );
	else
		myinet_putpeer(p);
}

static int mycleanup_once(unsigned long ttl)
{
	struct inet_peer *p;

	spin_lock_bh(&myinet_peer_unused_lock);
	p = myinet_peer_unused_head;
	if( p != NULL ){
		if( time_after(p->dtime + ttl, jiffies) ){
			spin_unlock_bh( &myinet_peer_unused_lock );
			return -1;
		}
		myinet_peer_unused_head = p->unused_next;
		if (p->unused_next != NULL)
			p->unused_next->unused_prevp = p->unused_prevp;
		else
			myinet_peer_unused_tailp = p->unused_prevp;
		p->unused_prevp = NULL;
		atomic_inc(&p->refcnt);
	}
	spin_unlock_bh( &myinet_peer_unused_lock );

	if (p == NULL)
		return -1;

	myunlink_from_pool(p);
	return 0;
}

struct inet_peer *myinet_getpeer(__u32 daddr, int create)
{
	struct inet_peer *p, *n;
	struct inet_peer **stack[PEER_MAXDEPTH], ***stackptr;

	read_lock_bh(&mypeer_pool_lock);
	p = mylookup(daddr);
	if( p != mypeer_avl_empty )
		atomic_inc( &p->refcnt );
	read_unlock_bh( &mypeer_pool_lock );

	if( p != mypeer_avl_empty ){
		myunlink_from_unused(p);
		return p;
	}

	if (!create)
		return NULL;

	n = kmem_cache_alloc( mypeer_cachep, GFP_ATOMIC );
	if (n == NULL)
		return NULL;
	printk(KERN_INFO "get peer: %u.%u.%u.%u\n", NIPQUAD(daddr) );
	n->v4daddr = daddr;
	atomic_set(&n->refcnt, 1);
	atomic_set(&n->rid, 0);
	//n->ip_id_count = secure_ip_id(daddr);
	n->tcp_ts_stamp = 0;

	write_lock_bh( &mypeer_pool_lock );
	p = mylookup(daddr);
	if( p != mypeer_avl_empty )
		goto out_free;

	mylink_to_pool( n );
	n->unused_prevp = NULL;
	mypeer_total++;
	write_unlock_bh( &mypeer_pool_lock );

	if( mypeer_total >= myinet_peer_threshold )
		mycleanup_once(0);
	return n;

out_free:
	atomic_inc(&p->refcnt);
	write_unlock_bh( &mypeer_pool_lock );
	myunlink_from_unused( p );
	kmem_cache_free( mypeer_cachep, n );
	return p;
}

void __init myinet_initpeers(void)
{
	struct sysinfo si;
	
	si_meminfo(&si);

	if (si.totalram <= (32768*1024)/PAGE_SIZE)
		myinet_peer_threshold >>= 1;
	if (si.totalram <= (16384*1024)/PAGE_SIZE)
		myinet_peer_threshold >>= 1;
	if (si.totalram <= (8192*1024)/PAGE_SIZE)
		myinet_peer_threshold >>= 2;

	mypeer_cachep = kmem_cache_create("myinet_peer_cache",
					sizeof(struct inet_peer),
					0, SLAB_HWCACHE_ALIGN,
					NULL, NULL);
	if (!mypeer_cachep)
		panic("cannot create inet_peer_cache");

	mypeer_periodic_timer.expires = jiffies + net_random() % myinet_peer_gc_maxtime
									+ myinet_peer_gc_maxtime;
	add_timer(&mypeer_periodic_timer);
}

static void __exit myinetpeers_exit(void)
{
	del_timer( &mypeer_periodic_timer );
	kmem_cache_destroy( mypeer_cachep );	
}